_METADATA

{"tree_metadata": {"('encoder', 'block', '0', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '0', 'SelfAttention', 'relative_attention_bias', 'embedding')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "relative_attention_bias", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '0', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '1', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '10', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "10", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '11', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "11", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '12', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "12", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '13', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "13", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '14', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "14", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '15', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "15", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '16', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "16", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '17', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "17", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '18', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "18", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '19', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "19", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '2', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '20', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "20", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '21', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "21", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '22', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "22", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '23', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "23", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '3', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '4', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '5', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '6', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '7', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "7", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '8', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "8", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '0', 'SelfAttention', 'k', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "k", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '0', 'SelfAttention', 'o', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "o", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '0', 'SelfAttention', 'q', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "q", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '0', 'SelfAttention', 'v', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "SelfAttention", "key_type": 2}, {"key": "v", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '0', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '1', 'DenseReluDense', 'wi_0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '1', 'DenseReluDense', 'wi_1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wi_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '1', 'DenseReluDense', 'wo', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "DenseReluDense", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'block', '9', 'layer', '1', 'layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "block", "key_type": 2}, {"key": "9", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('encoder', 'final_layer_norm', 'weight')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "final_layer_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}, "('shared', 'embedding')": {"key_metadata": [{"key": "shared", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}