siriuz42 commited on
Commit
3b5eae2
1 Parent(s): cb20a34

Upload initial checkpoint

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ checkpoints/checkpoint_1100000/state/checkpoint filter=lfs diff=lfs merge=lfs -text
checkpoints/checkpoint_1100000/descriptor/descriptor.pbtxt ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ uuid: "07e37595-56cb-444a-8f8e-abbc3c2e752f"
2
+ build_data {
3
+ timestamp {
4
+ seconds: 1714379165
5
+ }
6
+ user: "pax-dev-releaser-jobs"
7
+ hostname: "jgbv26.prod.google.com"
8
+ path: "/google/src/cloud/buildrabbit-username/buildrabbit-client/google3"
9
+ target: "//learning/multipod/pax/tools:colab_notebook"
10
+ invocation_id: "30f9e7fa-1cb7-4443-a49a-dbfb1b690df0"
11
+ changelist: 628988949
12
+ baseline_cl: 628988949
13
+ workspace_id: ""
14
+ client_status: BUILD_CLIENT_STATUS_UNSPECIFIED
15
+ verifiable: UNKNOWN
16
+ }
checkpoints/checkpoint_1100000/metadata/metadata ADDED
@@ -0,0 +1 @@
 
 
1
+ {"version": 1.1, "train_state_metadata": {"mdl_vars": {"params": {"freq_emb": {"emb_var": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [3, 1280]}}, "horizon_ff_layer": {"hidden_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "output_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "residual_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}}, "input_ff_layer": {"hidden_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [64, 1280]}}}, "output_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "residual_layer": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [64, 1280]}}}}, "stacked_transformer_layer": {"x_layers_0": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_1": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_10": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_11": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_12": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_13": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_14": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_15": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_16": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_17": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_18": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_19": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_2": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_3": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_4": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_5": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_6": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_7": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_8": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}, "x_layers_9": {"ff_layer": {"ffn_layer1": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "ffn_layer2": {"bias": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "linear": {"w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 1280]}}}, "layer_norm": {"bias": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}}, "layer_norm": {"scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}}, "self_attention": {"key": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "per_dim_scale": {"per_dim_scale": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [80]}}, "post": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "query": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}, "value": {"b": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [16, 80]}, "w": {"_array_metadata_tag": true, "dtype": "float32", "is_optax_masked_node": false, "unpadded_shape": [1280, 16, 80]}}}}}}}}}
checkpoints/checkpoint_1100000/state/checkpoint ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d147f95bbe40ffeb1e8d61643925dd29eeb00b5696db33bf9f6e1887682f6391
3
+ size 814291152