| Variable continuous_encoder/Embed_0/embedding size 196608 shape (vocab=256, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/encoder_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/input_proj/kernel size 98304 shape (vocab=128, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_0/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_0/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_0/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_0/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_0/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_0/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_0/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_0/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_0/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_1/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_1/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_1/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_1/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_1/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_1/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_1/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_1/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_1/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_10/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_10/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_10/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_10/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_10/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_10/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_10/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_10/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_10/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_11/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_11/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_11/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_11/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_11/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_11/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_11/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_11/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_11/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_2/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_2/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_2/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_2/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_2/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_2/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_2/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_2/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_2/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_3/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_3/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_3/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_3/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_3/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_3/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_3/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_3/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_3/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_4/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_4/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_4/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_4/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_4/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_4/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_4/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_4/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_4/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_5/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_5/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_5/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_5/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_5/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_5/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_5/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_5/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_5/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_6/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_6/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_6/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_6/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_6/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_6/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_6/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_6/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_6/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_7/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_7/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_7/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_7/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_7/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_7/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_7/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_7/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_7/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_8/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_8/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_8/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_8/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_8/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_8/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_8/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_8/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_8/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_9/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_9/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_9/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_9/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable continuous_encoder/layers_9/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_9/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable continuous_encoder/layers_9/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable continuous_encoder/layers_9/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable continuous_encoder/layers_9/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/Embed_0/embedding size 196608 shape (vocab=256, embed=768) partition spec ('model', None) |
| Variable decoder/continuous_inputs_projection/kernel size 98304 shape (vocab=128, embed=768) partition spec ('model', None) |
| Variable decoder/decoder_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_0/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_0/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_0/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_0/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_0/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_0/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_0/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_0/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_0/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_0/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_0/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_0/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_0/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_0/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_0/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_0/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_1/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_1/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_1/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_1/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_1/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_1/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_1/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_1/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_1/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_1/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_1/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_10/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_10/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_10/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_10/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_10/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_10/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_10/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_10/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_10/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_10/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_10/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_11/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_11/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_11/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_11/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_11/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_11/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_11/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_11/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_11/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_11/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_11/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_2/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_2/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_2/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_2/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_2/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_2/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_2/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_2/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_2/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_2/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_2/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_3/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_3/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_3/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_3/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_3/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_3/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_3/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_3/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_3/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_3/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_3/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_4/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_4/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_4/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_4/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_4/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_4/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_4/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_4/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_4/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_4/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_4/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_5/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_5/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_5/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_5/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_5/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_5/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_5/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_5/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_5/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_5/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_5/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_6/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_6/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_6/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_6/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_6/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_6/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_6/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_6/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_6/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_6/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_6/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_7/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_7/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_7/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_7/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_7/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_7/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_7/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_7/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_7/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_7/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_7/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_8/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_8/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_8/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_8/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_8/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_8/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_8/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_8/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_8/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_8/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_8/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/FiLMLayer_0/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_9/FiLMLayer_1/DenseGeneral_0/kernel size 4718592 shape (vocab=3072, embed=1536) partition spec ('model', None) |
| Variable decoder/layers_9/MultiHeadDotProductAttention_0/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/MultiHeadDotProductAttention_0/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_9/MultiHeadDotProductAttention_0/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/MultiHeadDotProductAttention_0/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_9/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable decoder/layers_9/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable decoder/layers_9/pre_cross_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_9/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_9/pre_self_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable decoder/layers_9/self_attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/self_attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable decoder/layers_9/self_attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/layers_9/self_attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable decoder/spec_out_dense/kernel size 98304 shape (embed=768, vocab=128) partition spec (None, 'model') |
| Variable decoder/time_emb_dense0/kernel size 2359296 shape (vocab=768, embed=3072) partition spec ('model', None) |
| Variable decoder/time_emb_dense1/kernel size 9437184 shape (vocab=3072, embed=3072) partition spec ('model', None) |
| Variable token_encoder/Embed_0/embedding size 1572864 shape (vocab=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/encoder_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_0/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_0/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_0/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_0/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_0/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_0/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_0/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_0/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_0/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_1/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_1/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_1/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_1/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_1/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_1/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_1/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_1/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_1/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_10/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_10/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_10/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_10/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_10/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_10/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_10/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_10/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_10/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_11/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_11/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_11/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_11/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_11/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_11/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_11/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_11/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_11/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_2/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_2/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_2/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_2/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_2/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_2/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_2/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_2/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_2/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_3/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_3/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_3/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_3/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_3/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_3/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_3/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_3/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_3/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_4/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_4/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_4/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_4/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_4/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_4/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_4/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_4/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_4/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_5/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_5/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_5/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_5/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_5/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_5/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_5/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_5/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_5/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_6/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_6/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_6/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_6/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_6/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_6/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_6/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_6/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_6/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_7/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_7/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_7/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_7/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_7/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_7/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_7/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_7/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_7/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_8/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_8/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_8/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_8/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_8/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_8/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_8/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_8/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_8/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_9/attention/key/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_9/attention/out/kernel size 589824 shape (joined_kv=768, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_9/attention/query/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_9/attention/value/kernel size 589824 shape (embed=768, joined_kv=768) partition spec (None, 'model') |
| Variable token_encoder/layers_9/mlp/wi_0/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_9/mlp/wi_1/kernel size 1572864 shape (embed=768, mlp=2048) partition spec (None, 'model') |
| Variable token_encoder/layers_9/mlp/wo/kernel size 1572864 shape (mlp=2048, embed=768) partition spec ('model', None) |
| Variable token_encoder/layers_9/pre_attention_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/layers_9/pre_mlp_layer_norm/scale size 768 shape (embed=768) partition spec (None,) |
| Variable token_encoder/token_embedder/embedding size 1179648 shape (vocab=1536, embed=768) partition spec ('model', None) |
| Total number of parameters: 411665664 |
|
|
| Variable param_states/continuous_encoder/Embed_0/embedding/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/Embed_0/embedding/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/Embed_0/embedding/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/Embed_0/embedding/v_row size 256 shape (256,) partition spec None |
| Variable param_states/continuous_encoder/encoder_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/encoder_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/encoder_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/encoder_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/input_proj/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/input_proj/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/input_proj/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/input_proj/kernel/v_row size 128 shape (128,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_0/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_0/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_0/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_1/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_1/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_1/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_10/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_10/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_10/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_11/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_11/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_11/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_2/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_2/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_2/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_3/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_3/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_3/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_4/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_4/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_4/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_5/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_5/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_5/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_6/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_6/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_6/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_7/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_7/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_7/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_8/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_8/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_8/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_9/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/continuous_encoder/layers_9/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/continuous_encoder/layers_9/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/Embed_0/embedding/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/Embed_0/embedding/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/Embed_0/embedding/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/Embed_0/embedding/v_row size 256 shape (256,) partition spec None |
| Variable param_states/decoder/continuous_inputs_projection/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/continuous_inputs_projection/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/continuous_inputs_projection/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/continuous_inputs_projection/kernel/v_row size 128 shape (128,) partition spec None |
| Variable param_states/decoder/decoder_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/decoder_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/decoder_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/decoder_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_0/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_0/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_0/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_0/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_0/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_0/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_1/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_1/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_1/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_1/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_1/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_1/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_10/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_10/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_10/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_10/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_10/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_10/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_11/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_11/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_11/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_11/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_11/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_11/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_2/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_2/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_2/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_2/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_2/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_2/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_3/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_3/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_3/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_3/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_3/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_3/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_4/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_4/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_4/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_4/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_4/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_4/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_5/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_5/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_5/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_5/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_5/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_5/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_6/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_6/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_6/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_6/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_6/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_6/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_7/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_7/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_7/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_7/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_7/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_7/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_8/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_8/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_8/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_8/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_8/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_8/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_0/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_0/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_0/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_0/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_1/DenseGeneral_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_1/DenseGeneral_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_1/DenseGeneral_0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/layers_9/FiLMLayer_1/DenseGeneral_0/kernel/v_row size 1536 shape (1536,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/MultiHeadDotProductAttention_0/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/decoder/layers_9/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/pre_cross_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_cross_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_9/pre_cross_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_cross_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_9/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_self_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_self_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/decoder/layers_9/pre_self_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/pre_self_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/layers_9/self_attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/spec_out_dense/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/spec_out_dense/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/spec_out_dense/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/decoder/spec_out_dense/kernel/v_row size 128 shape (128,) partition spec None |
| Variable param_states/decoder/time_emb_dense0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/time_emb_dense0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/time_emb_dense0/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/time_emb_dense0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/decoder/time_emb_dense1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/decoder/time_emb_dense1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/decoder/time_emb_dense1/kernel/v_col size 3072 shape (3072,) partition spec None |
| Variable param_states/decoder/time_emb_dense1/kernel/v_row size 3072 shape (3072,) partition spec None |
| Variable param_states/token_encoder/Embed_0/embedding/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/Embed_0/embedding/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/Embed_0/embedding/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/Embed_0/embedding/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/encoder_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/encoder_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/encoder_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/encoder_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_0/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_0/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_0/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_0/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_1/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_1/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_1/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_1/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_10/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_10/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_10/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_10/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_11/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_11/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_11/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_11/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_2/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_2/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_2/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_2/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_3/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_3/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_3/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_3/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_4/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_4/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_4/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_4/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_5/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_5/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_5/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_5/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_6/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_6/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_6/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_6/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_7/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_7/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_7/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_7/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_8/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_8/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_8/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_8/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/key/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/key/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/key/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/key/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/out/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/out/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/out/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/out/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/query/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/query/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/query/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/query/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/value/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/value/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/value/kernel/v_col size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/attention/value/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_0/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_0/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_0/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_0/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_1/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_1/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_1/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wi_1/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wo/kernel/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wo/kernel/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wo/kernel/v_col size 2048 shape (2048,) partition spec None |
| Variable param_states/token_encoder/layers_9/mlp/wo/kernel/v_row size 768 shape (768,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_attention_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_attention_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_9/pre_attention_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_attention_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_mlp_layer_norm/scale/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_mlp_layer_norm/scale/v size 768 shape (embed=768) partition spec (None,) |
| Variable param_states/token_encoder/layers_9/pre_mlp_layer_norm/scale/v_col size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/layers_9/pre_mlp_layer_norm/scale/v_row size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/token_embedder/embedding/m size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/token_embedder/embedding/v size 1 shape (1,) partition spec None |
| Variable param_states/token_encoder/token_embedder/embedding/v_col size 1536 shape (1536,) partition spec None |
| Variable param_states/token_encoder/token_embedder/embedding/v_row size 768 shape (768,) partition spec None |
| Variable step size 1 shape () partition spec None |
|
|
