Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,24 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/06820cf27a94ca47e9815a366132aa56 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/199039cc8b4b5f0df062476fdd3b8de1 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/29790f17f0b904f3d92d346dedff7505 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/6f8a7e02a7fb80d0193ca5709ed10897 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/760fb21414756632201e024ea2336fcd filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/894411cd0975d8d858842ec8a8018409 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/8b14e5e4f8e0827f34d11c8314ad58ae filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/94645838345dfdc099423302d7189ace filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/96fcddfede1497d4a90d59db5d8cb6c0 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/9b7e29d1f267772853b59d7169ff52dc filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/c33151b6c5295f0f25b159da7d851365 filter=lfs diff=lfs merge=lfs -text
+classification/orbax/1/params/ocdbt.process_0/d/d5aa581b19769b6c60f30f3c3cac968a filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/28fe402f3b097137b9c48fb0a3907676 filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/34d7c8f0ad6ab9d6254c853e56c2c2b2 filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/402c94ecedb1773139bc3f7521e85a57 filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/64b23608bd5e9e259f8514089efab6c1 filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/6de73bb54e278a8846bf06c51d6b24cf filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/72e7e06ea007cb5d5be66f04a3ea367e filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/7982a73b557f8e83c78b40830af24d4a filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/cc5fcf3244061d35b9285fd9633d38ad filter=lfs diff=lfs merge=lfs -text
+regression/orbax/1/params/ocdbt.process_0/d/ce7d920461e5fcd3ac1a1eb4c6ab3445 filter=lfs diff=lfs merge=lfs -text

classification/config.json

@@ -0,0 +1,6 @@
+{
+  "model_type": "tabfm",
+  "version": "1.0.0",
+  "task": "classification",
+  "framework": "flax/nnx"
+}

classification/orbax/1/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"params": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1782451148444418721, "commit_timestamp_nsecs": 1782451161317920056, "custom_metadata": {}}

classification/orbax/1/params/_METADATA

@@ -0,0 +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', 'embedding', 'value')": {"key_metadata": [{"key": "cell_embedder", "key_type": 2}, {"key": "y_embedder_lookup", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10, 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": [10]}}, "('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, 10]}}, "('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', 'projection', 'bias', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('icl_predictor', 'y_encoder', 'projection', 'kernel', 'value')": {"key_metadata": [{"key": "icl_predictor", "key_type": 2}, {"key": "y_encoder", "key_type": 2}, {"key": "projection", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10, 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}

classification/orbax/1/params/_sharding

@@ -0,0 +1 @@
+{"Y29sX2VtYmVkZGVyLm91dF93LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLm91dF93Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLmxuX3cuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MuaW5kX3ZlY3RvcnMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnBlcl9kaW1fc2NhbGUucGVyX2RpbV9zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5hdHRuLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5saW5lYXIxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5saW5lYXIxX2dhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5saW5lYXIyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5wb3N0X2F0dG5fbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5wb3N0X2ZmX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5wcmVfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMS5wcmVfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnBlcl9kaW1fc2NhbGUucGVyX2RpbV9zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5hdHRuLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5saW5lYXIxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5saW5lYXIxX2dhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5saW5lYXIyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5wb3N0X2F0dG5fbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5wb3N0X2ZmX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5wcmVfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyLnRmX2NvbC5ibG9ja3MubWFiMi5wcmVfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIub3V0X3cuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIub3V0X3cua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIubG5fdy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5pbmRfdmVjdG9ycy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ua19wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ua19wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ua2V5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ub3V0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ub3V0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ucGVyX2RpbV9zY2FsZS5wZXJfZGltX3NjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ucV9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ucV9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4ucXVlcnlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4udl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmF0dG4udl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmxpbmVhcjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmxpbmVhcjFfZ2F0ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLmxpbmVhcjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLnBvc3RfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLnBvc3RfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLnByZV9hdHRuX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIxLnByZV9mZl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ua19wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ua19wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ua2V5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ub3V0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ub3V0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ucGVyX2RpbV9zY2FsZS5wZXJfZGltX3NjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ucV9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ucV9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4ucXVlcnlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4udl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmF0dG4udl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmxpbmVhcjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmxpbmVhcjFfZ2F0ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLmxpbmVhcjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLnBvc3RfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLnBvc3RfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLnByZV9hdHRuX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y29sX2VtYmVkZGVyXzIudGZfY29sLmJsb2Nrcy5tYWIyLnByZV9mZl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci55X2VtYmVkZGVyX2xvb2t1cC5lbWJlZGRpbmcudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5mb3VyaWVyX2ZyZXF1ZW5jaWVzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5mb3VyaWVyX2ZyZXF1ZW5jaWVzX2NhdC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5pbl9saW5lYXIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5pbl9saW5lYXIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5pbl9saW5lYXJfY2F0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2VsbF9lbWJlZGRlci5pbl9saW5lYXJfY2F0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","Y2xzX3Rva2Vucy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ua19wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ua19wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ua2V5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ub3V0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ub3V0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ucGVyX2RpbV9zY2FsZS5wZXJfZGltX3NjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ucV9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ucV9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4ucXVlcnlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4udl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmF0dG4udl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmxpbmVhcjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmxpbmVhcjFfZ2F0ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLmxpbmVhcjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLnBvc3RfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLnBvc3RfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLnByZV9hdHRuX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci50Zl9pY2wuYmxvY2tzLnByZV9mZl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci55X2VuY29kZXIucHJvamVjdGlvbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci55X2VuY29kZXIucHJvamVjdGlvbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci5kZWNvZGVyLmxheWVycy4wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci5kZWNvZGVyLmxheWVycy4wLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci5kZWNvZGVyLmxheWVycy4yLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci5kZWNvZGVyLmxheWVycy4yLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","aWNsX3ByZWRpY3Rvci5sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3Iub3V0X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLmtleV9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnBlcl9kaW1fc2NhbGUucGVyX2RpbV9zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnF1ZXJ5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5hdHRuLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5saW5lYXIxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5saW5lYXIxX2dhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5saW5lYXIyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5wb3N0X2F0dG5fbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5wb3N0X2ZmX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5wcmVfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LmJsb2Nrcy5wcmVfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3IudGZfcm93LnJvcGUuZnJlcXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ua19wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ua19wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ua2V5X2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ub3V0X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ub3V0X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ucGVyX2RpbV9zY2FsZS5wZXJfZGltX3NjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ucV9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ucV9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4ucXVlcnlfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4udl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmF0dG4udl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmxpbmVhcjEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmxpbmVhcjFfZ2F0ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLmxpbmVhcjIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLnBvc3RfYXR0bl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLnBvc3RfZmZfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLnByZV9hdHRuX2xuLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cuYmxvY2tzLnByZV9mZl9sbi5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi50Zl9yb3cucm9wZS5mcmVxcy52YWx1ZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}","cm93X2ludGVyYWN0b3JfMi5vdXRfbG4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cpu:0\"}"}

classification/orbax/1/params/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "cell_embedder.fourier_frequencies.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.fourier_frequencies_cat.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.in_linear.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.in_linear.kernel.value", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.in_linear_cat.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.in_linear_cat.kernel.value", "write_shape": [64, 256], "chunk_shape": [64, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "cell_embedder.y_embedder_lookup.embedding.value", "write_shape": [10, 256], "chunk_shape": [10, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "cls_tokens.value", "write_shape": [8, 256], "chunk_shape": [8, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.ln_w.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.out_w.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.out_w.kernel.value", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.ind_vectors.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.key_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.query_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab1.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.key_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.query_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder.tf_col.blocks.mab2.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.ln_w.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.out_w.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.out_w.kernel.value", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.ind_vectors.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.key_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.query_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab1.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.key_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.query_ln.scale.value", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "col_embedder_2.tf_col.blocks.mab2.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.decoder.layers.0.bias.value", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.decoder.layers.0.kernel.value", "write_shape": [2048, 4096], "chunk_shape": [2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.decoder.layers.2.bias.value", "write_shape": [10], "chunk_shape": [10], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.decoder.layers.2.kernel.value", "write_shape": [4096, 10], "chunk_shape": [4096, 10], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.ln.scale.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.k_proj.bias.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.k_proj.kernel.value", "write_shape": [24, 2048, 2048], "chunk_shape": [24, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.key_ln.scale.value", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.out_proj.bias.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.out_proj.kernel.value", "write_shape": [24, 2048, 2048], "chunk_shape": [24, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.per_dim_scale.per_dim_scale.value", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.q_proj.bias.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.q_proj.kernel.value", "write_shape": [24, 2048, 2048], "chunk_shape": [24, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.query_ln.scale.value", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.v_proj.bias.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.attn.v_proj.kernel.value", "write_shape": [24, 2048, 2048], "chunk_shape": [24, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.linear1.kernel.value", "write_shape": [24, 2048, 8192], "chunk_shape": [24, 2048, 8192], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.linear1_gate.kernel.value", "write_shape": [24, 2048, 8192], "chunk_shape": [24, 2048, 8192], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.linear2.kernel.value", "write_shape": [24, 8192, 2048], "chunk_shape": [24, 8192, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.post_attn_ln.scale.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.post_ff_ln.scale.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.pre_attn_ln.scale.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.tf_icl.blocks.pre_ff_ln.scale.value", "write_shape": [24, 2048], "chunk_shape": [24, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.y_encoder.projection.bias.value", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "icl_predictor.y_encoder.projection.kernel.value", "write_shape": [10, 2048], "chunk_shape": [10, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.out_ln.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.key_ln.scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.query_ln.scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.blocks.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor.tf_row.rope.freqs.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.out_ln.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.k_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.k_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.key_ln.scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.out_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.out_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.per_dim_scale.per_dim_scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.q_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.q_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.query_ln.scale.value", "write_shape": [3, 32], "chunk_shape": [3, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.v_proj.bias.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.attn.v_proj.kernel.value", "write_shape": [3, 256, 256], "chunk_shape": [3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.linear1.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.linear1_gate.kernel.value", "write_shape": [3, 256, 1024], "chunk_shape": [3, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.linear2.kernel.value", "write_shape": [3, 1024, 256], "chunk_shape": [3, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.post_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.post_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.pre_attn_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.blocks.pre_ff_ln.scale.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "row_interactor_2.tf_row.rope.freqs.value", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}]}

classification/orbax/1/params/descriptor/descriptor.pbtxt

@@ -0,0 +1,8 @@
+uuid: "06c0943b-3b73-4bcb-82a1-8253e64bee2c"
+build_data {
+  baseline_cl: 934720893
+  workspace_id: "weihaokong/740"
+  citc_snapshot: 10966
+  client_status: BUILD_CLIENT_STATUS_UNSPECIFIED
+  verifiable: UNKNOWN
+}

classification/orbax/1/params/descriptor/uuid-06c0943b-3b73-4bcb-82a1-8253e64bee2c

@@ -0,0 +1 @@
+uuid: "06c0943b-3b73-4bcb-82a1-8253e64bee2c"

classification/orbax/1/params/ocdbt.process_0/d/06820cf27a94ca47e9815a366132aa56

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2b9b45c3f694f9074202029ff8c50b18f53c69ab36534ab22ae0623a7cd47b9
+size 4399104

classification/orbax/1/params/ocdbt.process_0/d/199039cc8b4b5f0df062476fdd3b8de1

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad54b6d12f00ade6b26cb1058b215b79d90863db6ba0ffaba4d04e54aa6dcf9d
+size 373497856

classification/orbax/1/params/ocdbt.process_0/d/29790f17f0b904f3d92d346dedff7505

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:907974a532d8e8d395c59e41cab3490dff5a11a826a55c2ab7791405e877205b
+size 2992283648