diff --git "a/checkpoints/10/default/_METADATA" "b/checkpoints/10/default/_METADATA"
--- "a/checkpoints/10/default/_METADATA"
+++ "b/checkpoints/10/default/_METADATA"
@@ -1 +1 @@
-{"tree_metadata": {"('model', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '0', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '0', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '1', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '1', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '2', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '2', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '3', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '3', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '4', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '4', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '5', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '5', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('model', 'backbone', 'norm', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'patch_embed', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('model', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('model', 'classifier_head', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('model', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('model', 'grounding', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'grounding', 'projection', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('model', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('model', 'logic', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', '_pytree__state')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '0')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '1', '0', 'count', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '1')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '1', '2', 'count', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'step', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "step", "key_type": 2}, {"key": "value", "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
+{"tree_metadata": {"('model', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '0', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '0', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '0', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '0', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '1', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '1', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '1', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '1', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '2', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '2', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '2', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '2', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '3', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '3', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '3', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '3', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '4', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '4', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '4', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '4', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('model', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('model', 'backbone', 'blocks', '5', 'dropout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', 'count', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'backbone', 'blocks', '5', 'dropout', 'rngs', 'key', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "dropout", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'backbone', 'blocks', '5', 'mlp', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('model', 'backbone', 'blocks', '5', 'mlp', 'layers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'blocks', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('model', 'backbone', 'norm', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'patch_embed', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('model', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('model', 'classifier_head', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('model', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('model', 'grounding', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'grounding', 'bn', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'grounding', 'bn', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'grounding', 'bn', 'mean', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "mean", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'grounding', 'bn', 'scale', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'grounding', 'bn', 'var', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "var", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('model', 'grounding', 'projection', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('model', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('model', 'logic', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '0', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '0', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '1', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '1', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '2', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '2', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '3', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '3', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '4', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '4', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '5', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '5', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '6', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '6', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '7', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '7', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '8', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '8', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '9', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '9', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '10', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '10', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '11', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '11', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '12', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '12', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '13', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '13', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'child', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "child", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '1', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'children', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('model', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('model', 'logic', 'rules', '14', 'root', 'right', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '14', 'root', 'right', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 1}, {"key": "root", "key_type": 2}, {"key": "right", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('model', 'logic', 'rules', '_pytree__state')": {"key_metadata": [{"key": "model", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', '_pytree__state')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "_pytree__state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '0')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '1', '0', 'count', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'bn', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'bn', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'mu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '0', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '1', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '2', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '3', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '4', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_L', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_L", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'proj_U', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj_U", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'attn', 'qkv', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1152]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 1536]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'mlp', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm1', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm1', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm2', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'blocks', '5', 'norm2', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'cls_token', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "cls_token", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'norm', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'norm', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'patch_embed', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'patch_embed', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 16, 3, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'backbone', 'pos_embed', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "pos_embed", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 197, 384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'classifier_head', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'classifier_head', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "classifier_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47, 10]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'bn', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'bn', 'scale', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "bn", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [47]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'grounding', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "grounding", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 47]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'body', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "body", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'head', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'predicates', 'mouth', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mouth", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '0', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'boxy', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boxy", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'metallic', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "metallic", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'structure', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "structure", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'predicates', 'windows', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "windows", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '1', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'is_cat', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_cat", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'small_ears', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "small_ears", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '2', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'beak', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "beak", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'feathers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "feathers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'is_bird', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_bird", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'two_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "two_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '3', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'antlers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "antlers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'is_deer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_deer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'predicates', 'short_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "short_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '4', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'is_dog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_dog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'long_tail', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_tail", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'longer_snout', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_snout", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'predicates', 'whiskers', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "whiskers", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '5', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'bulging_eyes', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "bulging_eyes", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'colorful', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "colorful", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'is_frog', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_frog", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'predicates', 'shiny', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "shiny", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '6', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'four_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "four_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'hooves', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "hooves", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_animal', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_animal", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'is_horse', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_horse", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'long_legs', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "long_legs", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'longer_neck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer_neck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'predicates', 'mane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "mane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '7', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_automobile', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_automobile", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '8', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'propeller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "propeller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '9', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_airplane', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_airplane", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'jet_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "jet_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '10', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'sails', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "sails", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '11', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'boat_engine', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "boat_engine", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '12', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'chimney', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "chimney", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_ship', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_ship", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'water', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "water", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'predicates', 'wings', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wings", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '13', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_truck', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_truck", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'is_vehicle', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "is_vehicle", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'longer', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "longer", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'smaller', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "smaller", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'offset_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "offset_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_l', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_l", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'predicates', 'wheels', 'predicate', 'slope_u', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "predicates", "key_type": 2}, {"key": "wheels", "key_type": 2}, {"key": "predicate", "key_type": 2}, {"key": "slope_u", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'children', '1', 'gate', 'weight', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'children', '0', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "children", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'beta', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "beta", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'opt_state', '1', '0', 'nu', 'logic', 'rules', '14', 'root', 'left', 'gate', 'weights', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "logic", "key_type": 2}, {"key": "rules", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "root", "key_type": 2}, {"key": "left", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('optimizer', 'opt_state', '1', '1')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('optimizer', 'opt_state', '1', '2', 'count', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}, {"key": "count", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('optimizer', 'step', 'value')": {"key_metadata": [{"key": "optimizer", "key_type": 2}, {"key": "step", "key_type": 2}, {"key": "value", "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