| 72 | } |
| 73 | |
| 74 | static const char * label_for_category(const std::string & category) { |
| 75 | if (category == "token_ids") return "Token IDs"; |
| 76 | if (category == "causal_mask") return "Causal mask"; |
| 77 | if (category == "text_token_embed") return "Token embedding"; |
| 78 | if (category == "text_after_pos_embed") return "Embedding + position"; |
| 79 | if (category == "block_after_ln1") return "Blocks after ln1"; |
| 80 | if (category == "block_qkv") return "Blocks qkv"; |
| 81 | if (category == "block_attn_out") return "Blocks attn out"; |
| 82 | if (category == "block_after_attn_residual") return "Blocks attn residual"; |
| 83 | if (category == "block_after_ln2") return "Blocks after ln2"; |
| 84 | if (category == "block_mlp_fc1") return "Blocks mlp fc1"; |
| 85 | if (category == "block_mlp_gelu") return "Blocks mlp gelu"; |
| 86 | if (category == "block_mlp_out") return "Blocks mlp out"; |
| 87 | if (category == "block_out") return "Blocks output"; |
| 88 | if (category == "text_final_ln") return "Final layer norm"; |
| 89 | if (category == "text_features_2d") return "Final text features"; |
| 90 | return category.c_str(); |
| 91 | } |
| 92 | |
| 93 | static std::vector<std::string> tensor_names() { |
| 94 | std::vector<std::string> names = { |