regression/orbax/1/params/_METADATA
| 1 | {"tree_metadata": {"('cell_embedder', 'fourier_frequencies', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "fourier_frequencies", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('cell_embedder', 'fourier_frequencies_cat', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "fourier_frequencies_cat", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('cell_embedder', 'in_linear', 'bias', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "in_linear", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('cell_embedder', 'in_linear', 'kernel', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "in_linear", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('cell_embedder', 'in_linear_cat', 'bias', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "in_linear_cat", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('cell_embedder', 'in_linear_cat', 'kernel', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "in_linear_cat", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 256]}}, "('cell_embedder', 'y_embedder_lookup', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "y_embedder_lookup", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6]}}, "('cell_embedder', 'y_embedder_lookup', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "y_embedder_lookup", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 6]}}, "('cell_embedder', 'y_embedder_lookup', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "y_embedder_lookup", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('cell_embedder', 'y_embedder_lookup', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "y_embedder_lookup", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 256]}}, "('cls_tokens', 'value')": {"key_metadata": [{"key": "cls_tokens", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 256]}}, "('col_embedder', 'ln_w', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "ln_w", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('col_embedder', 'out_w', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "out_w", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('col_embedder', 'out_w', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "out_w", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'ind_vectors', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ind_vectors", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab1', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder', 'tf_col', 'blocks', 'mab2', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'ln_w', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "ln_w", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('col_embedder_2', 'out_w', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "out_w", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('col_embedder_2', 'out_w', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "out_w", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'ind_vectors', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ind_vectors", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab1', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab1", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('col_embedder_2', 'tf_col', 'blocks', 'mab2', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "col_embedder_2", "key_type": 2}, {"key": "tf_col", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "mab2", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('icl_predictor', 'decoder', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('icl_predictor', 'decoder', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 4096]}}, "('icl_predictor', 'decoder', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('icl_predictor', 'decoder', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1]}}, "('icl_predictor', 'ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 8192]}}, "('icl_predictor', 'tf_icl', 'blocks', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 8192]}}, "('icl_predictor', 'tf_icl', 'blocks', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 8192, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'tf_icl', 'blocks', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "tf_icl", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048]}}, "('icl_predictor', 'y_encoder', 'layers', '0', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('icl_predictor', 'y_encoder', 'layers', '0', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 4096]}}, "('icl_predictor', 'y_encoder', 'layers', '2', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('icl_predictor', 'y_encoder', 'layers', '2', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 2048]}}, "('row_interactor', 'out_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "out_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('row_interactor', 'tf_row', 'blocks', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('row_interactor', 'tf_row', 'blocks', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'blocks', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor', 'tf_row', 'rope', 'freqs', 'value')": {"key_metadata": [{"key": "row_interactor", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "rope", "key_type": 2}, {"key": "freqs", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('row_interactor_2', 'out_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "out_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'key_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'per_dim_scale', 'per_dim_scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "per_dim_scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'query_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'linear1', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('row_interactor_2', 'tf_row', 'blocks', 'linear1_gate', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear1_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256, 1024]}}, "('row_interactor_2', 'tf_row', 'blocks', 'linear2', 'kernel', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "linear2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'post_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'post_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "post_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'pre_attn_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_attn_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'blocks', 'pre_ff_ln', 'scale', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "pre_ff_ln", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('row_interactor_2', 'tf_row', 'rope', 'freqs', 'value')": {"key_metadata": [{"key": "row_interactor_2", "key_type": 2}, {"key": "tf_row", "key_type": 2}, {"key": "rope", "key_type": 2}, {"key": "freqs", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} |