FuSe_PaliVLA_model/15000/model_model_params/_METADATA

{"tree_metadata": {"('params', 'digit_left', 'block0', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block0', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block1', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block10', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block11', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block2', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block3', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block4', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block5', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block6', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block7', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block8', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'block9', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'cls_token')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "cls_token", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'fc_norm', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "fc_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'fc_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "fc_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'head', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'head', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'patch_embed', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'patch_embed', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'attention_head', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pool_func', 'readouts')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "readouts", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_left', 'pos_embed')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_left", "key_type": 2}, {"key": "pos_embed", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block0', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block0", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block1', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block10', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block10", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block11', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block11", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block2', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block3', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block4', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block5', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block5", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block6', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block6", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block7', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block7", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block8', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block8", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'attn', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'attn', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'attn', 'qkv', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'attn', 'qkv', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "qkv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'mlp', 'fc1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'mlp', 'fc1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'mlp', 'fc2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'mlp', 'fc2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'norm1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'norm1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'norm2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'block9', 'norm2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "block9", "key_type": 2}, {"key": "norm2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'cls_token')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "cls_token", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'fc_norm', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "fc_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'fc_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "fc_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'head', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'head', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'patch_embed', 'proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'patch_embed', 'proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "patch_embed", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'attention_head', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "attention_head", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pool_func', 'readouts')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pool_func", "key_type": 2}, {"key": "readouts", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'digit_right', 'pos_embed')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "digit_right", "key_type": 2}, {"key": "pos_embed", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoder_norm', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoder_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'embedding', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'embedding', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'head', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'head', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'img', 'pos_embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'embedder', 'input_embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'final_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'attn', 'kv_einsum', 'w')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'attn', 'q_einsum', 'w')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'mlp', 'gating_einsum')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'mlp', 'linear')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'pre_attention_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'llm', 'layers', 'pre_ffw_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'modality_embedder', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "modality_embedder", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'conv_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit1', 'gn_proj', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block1', 'unit2', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block1", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'conv_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit1', 'gn_proj', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block2', 'unit2', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block2", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'conv_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit1', 'gn_proj', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block3', 'unit2', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block3", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'conv_proj', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "conv_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn_proj', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit1', 'gn_proj', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit1", "key_type": 2}, {"key": "gn_proj", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'conv1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'conv2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'conv3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "conv3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn2', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn2', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn3', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'block4', 'unit2', 'gn3', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "block4", "key_type": 2}, {"key": "unit2", "key_type": 2}, {"key": "gn3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'conv_root', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "conv_root", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'gn_root', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "gn_root", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'spectro', 'gn_root', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "spectro", "key_type": 2}, {"key": "gn_root", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_image_digit_left')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_image_digit_left", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_image_digit_right')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_image_digit_right", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_image_primary')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_image_primary", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_image_wrist')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_image_wrist", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_mel_spectro')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_mel_spectro", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_modality_idx')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_modality_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('params', 'start_text')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "start_text", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}}, "use_zarr3": false}