std::stringstream ss;

ss << "System Info: \n";

ss << " BLAS = " << ggml_cpu_has_blas() << std::endl;

ss << " SSE3 = " << ggml_cpu_has_sse3() << std::endl;

ss << " AVX = " << ggml_cpu_has_avx() << std::endl;

ss << " AVX2 = " << ggml_cpu_has_avx2() << std::endl;

ss << " AVX512 = " << ggml_cpu_has_avx512() << std::endl;

ss << " AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << std::endl;

ss << " AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << std::endl;

ss << " FMA = " << ggml_cpu_has_fma() << std::endl;

ss << " NEON = " << ggml_cpu_has_neon() << std::endl;

ss << " ARM_FMA = " << ggml_cpu_has_arm_fma() << std::endl;

ss << " F16C = " << ggml_cpu_has_f16c() << std::endl;

ss << " FP16_VA = " << ggml_cpu_has_fp16_va() << std::endl;

ss << " WASM_SIMD = " << ggml_cpu_has_wasm_simd() << std::endl;

ss << " VSX = " << ggml_cpu_has_vsx() << std::endl;

return ss.str();

uint32_t n = (uint32_t)ggml_nelements(tensor);

std::vector<float> random_numbers = rng->randn(n);

for (uint32_t i = 0; i < n; i++) {

ggml_set_f32_1d(tensor, i, random_numbers[i]);

}

std::string progress = " |";

int max_progress = 50;

int32_t current = (int32_t)(step * 1.f * max_progress / steps);

for (int i = 0; i < 50; i++) {

if (i > current) {

progress += " ";

} else if (i == current && i != max_progress - 1) {

progress += ">";

} else {

progress += "=";

}

}

progress += "|";

printf(time > 1.0f ? "\r%s %i/%i - %.2fs/it" : "\r%s %i/%i - %.2fit/s",

progress.c_str(), step, steps,

time > 1.0f || time == 0 ? time : (1.0f / time));

fflush(stdout); // for linux

if (step == steps) {

printf("\n");

}

printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);

fflush(stdout);

if (shape_only) {

return;

}

int range = 3;

for (int i = 0; i < tensor->ne[3]; i++) {

if (i >= range && i + range < tensor->ne[3]) {

continue;

}

for (int j = 0; j < tensor->ne[2]; j++) {

if (j >= range && j + range < tensor->ne[2]) {

continue;

}

for (int k = 0; k < tensor->ne[1]; k++) {

if (k >= range && k + range < tensor->ne[1]) {

continue;

}

for (int l = 0; l < tensor->ne[0]; l++) {

if (l >= range && l + range < tensor->ne[0]) {

continue;

}

if (tensor->type == GGML_TYPE_F32) {

printf(" [%d, %d, %d, %d] = %f\n", i, j, k, l, ggml_tensor_get_f32(tensor, l, k, j, i));

} else if (tensor->type == GGML_TYPE_F16) {

printf(" [%d, %d, %d, %d] = %i\n", i, j, k, l, ggml_tensor_get_f16(tensor, l, k, j, i));

}

fflush(stdout);

}

}

}

}

size_t ret = std::fread(ptr, size, count, stream);

if (ret != count) {

printf("Error: read from file failed");

exit(1);

}

if (dst->type == src->type) {

dst->nb[0] = src->nb[0];

dst->nb[1] = src->nb[1];

dst->nb[2] = src->nb[2];

dst->nb[3] = src->nb[3];

memcpy(((char*)dst->data), ((char*)src->data), ggml_nbytes(dst));

return;

}

struct ggml_init_params params;

params.mem_size = 10 * 1024 * 1024; // for padding

params.mem_buffer = NULL;

params.no_alloc = false;

struct ggml_context* ctx = ggml_init(params);

if (!ctx) {

LOG_ERROR("ggml_init() failed");

return;

}

ggml_tensor* final = ggml_cpy_inplace(ctx, src, dst);

struct ggml_cgraph* graph = ggml_new_graph(ctx);

ggml_build_forward_expand(graph, final);

ggml_graph_compute_with_ctx(ctx, graph, 1);

ggml_free(ctx);

float linear_start = 0.00085f,

float linear_end = 0.0120,

int timesteps = TIMESTEPS) {

float ls_sqrt = sqrtf(linear_start);

float le_sqrt = sqrtf(linear_end);

float amount = le_sqrt - ls_sqrt;

float product = 1.0f;

for (int i = 0; i < timesteps; i++) {

float beta = ls_sqrt + amount * ((float)i / (timesteps - 1));

product *= 1.0f - powf(beta, 2.0f);

alphas_cumprod[i] = product;

}

// timesteps: [N,]

// embedding: [(dim + 1)/2, N]

int half = dim / 2;

std::vector<float> freqs(half);

for (int i = 0; i < half; ++i) {

freqs[i] = (float)std::exp(-std::log(max_period) * i / half);

}

for (int i = 0; i < timesteps->ne[0]; ++i) {

for (int j = 0; j < half; ++j) {

float arg = ggml_get_f32_1d(timesteps, i) * freqs[j];

ggml_tensor_set_f32(embedding, std::cos(arg), j, i);

ggml_tensor_set_f32(embedding, std::sin(arg), j + half, i);

}

if (dim % 2 != 0) {

*(float*)((char*)embedding->data + i * embedding->nb[1] + dim * embedding->nb[0]) = 0;

}

}

struct ggml_tensor* output) {

int64_t width = output->ne[0];

int64_t height = output->ne[1];

int64_t channels = output->ne[2];

GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);

for (int iy = 0; iy < height; iy++) {

for (int ix = 0; ix < width; ix++) {

for (int k = 0; k < channels; k++) {

float value = *(image_data + iy * width * channels + ix * channels + k);

ggml_tensor_set_f32(output, value / 255.0f, ix, iy, k);

}

}

}

float mean = 0.0f;

int64_t nelements = ggml_nelements(src);

float* data = (float*)src->data;

for (int i = 0; i < nelements; i++) {

mean += data[i] / nelements * 1.0f;

}

return mean;

GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));

int64_t nelements = ggml_nelements(a);

float* vec_a = (float*)a->data;

float* vec_b = (float*)b->data;

for (int i = 0; i < nelements; i++) {

vec_a[i] = vec_a[i] + vec_b[i];

}

int64_t nelements = ggml_nelements(src);

float* data = (float*)src->data;

for (int i = 0; i < nelements; i++) {

data[i] = data[i] * scale;

}

int64_t nelements = ggml_nelements(src);

float* data = (float*)src->data;

for (int i = 0; i < nelements; i++) {

float val = data[i];

data[i] = val < min ? min : (val > max ? max : val);

}

int64_t nelements = ggml_nelements(src);

float* data = (float*)src->data;

for (int i = 0; i < nelements; i++) {

float val = data[i];

data[i] = val * 2.0f - 1.0f;

}

int64_t nelements = ggml_nelements(src);

float* data = (float*)src->data;

for (int i = 0; i < nelements; i++) {

float val = data[i];

data[i] = (val + 1.0f) * 0.5f;

}

std::regex re("<lora:([^:]+):([^>]+)>");

std::smatch matches;

std::unordered_map<std::string, float> filename2multiplier;

while (std::regex_search(text, matches, re)) {

std::string filename = matches[1].str();

float multiplier = std::stof(matches[2].str());

text = std::regex_replace(text, re, "", std::regex_constants::format_first_only);

if (multiplier == 0.f) {

continue;

}

if (filename2multiplier.find(filename) == filename2multiplier.end()) {

filename2multiplier[filename] = multiplier;

} else {

filename2multiplier[filename] += multiplier;

}

}

return std::make_pair(filename2multiplier, text);

std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;

std::set<int> byte_set;

for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {

byte_set.insert(b);

byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));

}

for (int b = 161; b <= 172; ++b) {

byte_set.insert(b);

byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));

}

for (int b = 174; b <= 255; ++b) {

byte_set.insert(b);

byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));

}

int n = 0;

for (int b = 0; b < 256; ++b) {

if (byte_set.find(b) == byte_set.end()) {

byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));

++n;

}

}

// LOG_DEBUG("byte_unicode_pairs %d", byte_unicode_pairs.size());

return byte_unicode_pairs;

SDVersion version = VERSION_1_x;

std::map<int, std::u32string> byte_encoder;

std::map<std::u32string, int> encoder;

std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;

std::regex pat;

static std::string strip(const std::string& str) {

std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");

std::string::size_type end = str.find_last_not_of(" \t\n\r\v\f");

if (start == std::string::npos) {

// String contains only whitespace characters

return "";

}

return str.substr(start, end - start + 1);

}

static std::string whitespace_clean(std::string text) {

text = std::regex_replace(text, std::regex(R"(\s+)"), " ");

text = strip(text);

return text;

}

static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {

std::set<std::pair<std::u32string, std::u32string>> pairs;

if (subwords.size() == 0) {

return pairs;

}

std::u32string prev_subword = subwords[0];

for (int i = 1; i < subwords.size(); i++) {

std::u32string subword = subwords[i];

std::pair<std::u32string, std::u32string> pair(prev_subword, subword);

pairs.insert(pair);

prev_subword = subword;

}

return pairs;

}

CLIPTokenizer(SDVersion version = VERSION_1_x)

: version(version) {}

void load_from_merges(const std::string& merges_utf8_str) {

auto byte_unicode_pairs = bytes_to_unicode();

byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());

// for (auto & pair: byte_unicode_pairs) {

// std::cout << pair.first << ": " << pair.second << std::endl;

// }

std::vector<std::u32string> merges;

size_t start = 0;

size_t pos;

std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);

while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {

merges.push_back(merges_utf32_str.substr(start, pos - start));

start = pos + 1;

}

// LOG_DEBUG("merges size %llu", merges.size());

GGML_ASSERT(merges.size() == 48895);

merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());

std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;

for (const auto& merge : merges) {

size_t space_pos = merge.find(' ');

merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));

// LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());

}

std::vector<std::u32string> vocab;

for (const auto& pair : byte_unicode_pairs) {

vocab.push_back(pair.second);

}

for (const auto& pair : byte_unicode_pairs) {

vocab.push_back(pair.second + utf8_to_utf32("</w>"));

}

for (const auto& merge : merge_pairs) {

vocab.push_back(merge.first + merge.second);

}

vocab.push_back(utf8_to_utf32("<|startoftext|>"));

vocab.push_back(utf8_to_utf32("<|endoftext|>"));

LOG_DEBUG("vocab size: %llu", vocab.size());

int i = 0;

for (const auto& token : vocab) {

encoder[token] = i++;

}

int rank = 0;

for (const auto& merge : merge_pairs) {

bpe_ranks[merge] = rank++;

}

};

std::u32string bpe(const std::u32string& token) {

std::vector<std::u32string> word;

for (int i = 0; i < token.size() - 1; i++) {

word.emplace_back(1, token[i]);

}

word.push_back(token.substr(token.size() - 1) + utf8_to_utf32("</w>"));

std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);

if (pairs.empty()) {

return token + utf8_to_utf32("</w>");

}

while (true) {

auto min_pair_iter = std::min_element(pairs.begin(),

pairs.end(),

[&](const std::pair<std::u32string, std::u32string>& a,

const std::pair<std::u32string, std::u32string>& b) {

if (bpe_ranks.find(a) == bpe_ranks.end()) {

return false;

} else if (bpe_ranks.find(b) == bpe_ranks.end()) {

return true;

}

return bpe_ranks.at(a) < bpe_ranks.at(b);

});

const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;

if (bpe_ranks.find(bigram) == bpe_ranks.end()) {

break;

}

std::u32string first = bigram.first;

std::u32string second = bigram.second;

std::vector<std::u32string> new_word;

int32_t i = 0;

while (i < word.size()) {

auto it = std::find(word.begin() + i, word.end(), first);

if (it == word.end()) {

new_word.insert(new_word.end(), word.begin() + i, word.end());

break;

}

new_word.insert(new_word.end(), word.begin() + i, it);

i = static_cast<int32_t>(std::distance(word.begin(), it));

if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {

new_word.push_back(first + second);

i += 2;

} else {

new_word.push_back(word[i]);

i += 1;

}

}

word = new_word;

if (word.size() == 1) {

break;

}

pairs = get_pairs(word);

}

std::u32string result;

for (int i = 0; i < word.size(); i++) {

result += word[i];

if (i != word.size() - 1) {

result += utf8_to_utf32(" ");

}

}

return result;

}

std::vector<int> tokenize(std::string text, size_t max_length = 0, bool padding = false) {

std::vector<int32_t> tokens = encode(text);

tokens.insert(tokens.begin(), BOS_TOKEN_ID);

if (max_length > 0) {

if (tokens.size() > max_length - 1) {

tokens.resize(max_length - 1);

tokens.push_back(EOS_TOKEN_ID);

} else {

tokens.push_back(EOS_TOKEN_ID);

if (padding) {

int pad_token_id = PAD_TOKEN_ID;

if (version == VERSION_2_x) {

pad_token_id = 0;

}

tokens.insert(tokens.end(), max_length - tokens.size(), pad_token_id);

}

}

}

return tokens;

}

std::vector<int> encode(std::string text) {

std::string original_text = text;

std::vector<int32_t> bpe_tokens;

text = whitespace_clean(text);

std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return std::tolower(c); });

std::regex pat(R"(<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",

std::regex::icase);

std::smatch matches;

std::string str = text;

std::vector<std::string> token_strs;

while (std::regex_search(str, matches, pat)) {

for (auto& token : matches) {

std::string token_str = token.str();

std::u32string utf32_token;

for (int i = 0; i < token_str.length(); i++) {

char b = token_str[i];

utf32_token += byte_encoder[b];

}

auto bpe_strs = bpe(utf32_token);

size_t start = 0;

size_t pos;

while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {

auto bpe_str = bpe_strs.substr(start, pos - start);

bpe_tokens.push_back(encoder[bpe_str]);

token_strs.push_back(utf32_to_utf8(bpe_str));

start = pos + 1;

}

auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);

bpe_tokens.push_back(encoder[bpe_str]);

token_strs.push_back(utf32_to_utf8(bpe_str));

}

str = matches.suffix();

}

std::stringstream ss;

ss << "[";

for (auto token : token_strs) {

ss << "\"" << token << "\", ";

}

ss << "]";

LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());

return bpe_tokens;

}

std::vector<std::pair<std::string, float>> res;

std::vector<int> round_brackets;

std::vector<int> square_brackets;

float round_bracket_multiplier = 1.1f;

float square_bracket_multiplier = 1 / 1.1f;

std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|\)|\]|[^\\()\[\]:]+|:)");

std::regex re_break(R"(\s*\bBREAK\b\s*)");

auto multiply_range = [&](int start_position, float multiplier) {

for (int p = start_position; p < res.size(); ++p) {

res[p].second *= multiplier;

}

};

std::smatch m;

std::string remaining_text = text;

while (std::regex_search(remaining_text, m, re_attention)) {

std::string text = m[0];

std::string weight = m[1];

if (text == "(") {

round_brackets.push_back((int)res.size());

} else if (text == "[") {

square_brackets.push_back((int)res.size());

} else if (!weight.empty()) {

if (!round_brackets.empty()) {

multiply_range(round_brackets.back(), std::stof(weight));

round_brackets.pop_back();

}

} else if (text == ")" && !round_brackets.empty()) {

multiply_range(round_brackets.back(), round_bracket_multiplier);

round_brackets.pop_back();

} else if (text == "]" && !square_brackets.empty()) {

multiply_range(square_brackets.back(), square_bracket_multiplier);

square_brackets.pop_back();

} else if (text == "\\(") {

res.push_back({text.substr(1), 1.0f});

} else {

res.push_back({text, 1.0f});

}

remaining_text = m.suffix();

}

for (int pos : round_brackets) {

multiply_range(pos, round_bracket_multiplier);

}

for (int pos : square_brackets) {

multiply_range(pos, square_bracket_multiplier);

}

if (res.empty()) {

res.push_back({"", 1.0f});

}

int i = 0;

while (i + 1 < res.size()) {

if (res[i].second == res[i + 1].second) {

res[i].first += res[i + 1].first;

res.erase(res.begin() + i + 1);

} else {

++i;

}

}

return res;

int32_t n_head;

int32_t d_model;

int32_t hidden_size; // n_head * d_model

int32_t intermediate_size;

// attention

struct ggml_tensor* q_w; // [hidden_size, hidden_size]

struct ggml_tensor* q_b; // [hidden_size, ]

struct ggml_tensor* k_w; // [hidden_size, hidden_size]

struct ggml_tensor* k_b; // [hidden_size, ]

struct ggml_tensor* v_w; // [hidden_size, hidden_size]

struct ggml_tensor* v_b; // [hidden_size, ]

struct ggml_tensor* out_w; // [hidden_size, hidden_size]

struct ggml_tensor* out_b; // [hidden_size, ]

// layer norm 1

struct ggml_tensor* ln1_w; // [hidden_size, ]

struct ggml_tensor* ln1_b; // [hidden_size, ]

// mlp

struct ggml_tensor* fc1_w; // [intermediate_size, hidden_size]

struct ggml_tensor* fc1_b; // [intermediate_size, ]

struct ggml_tensor* fc2_w; // [hidden_size, intermediate_size]

struct ggml_tensor* fc2_b; // [hidden_size, ]

// layer norm 2

struct ggml_tensor* ln2_w; // [hidden_size, ]

struct ggml_tensor* ln2_b; // [hidden_size, ]

struct ggml_tensor* attn_scale; // [hidden_size, ]

size_t calculate_mem_size(ggml_type wtype) {

double mem_size = 0;

mem_size += 4 * hidden_size * hidden_size * ggml_type_sizef(wtype); // q_w/k_w/v_w/out_w

mem_size += 8 * hidden_size * ggml_type_sizef(GGML_TYPE_F32); // q_b/k_b/v_b/out_b/ln1_w/ln1_b/ln2_w/ln2_b

mem_size += 2 * hidden_size * intermediate_size * ggml_type_sizef(wtype); // fc1_w/fc2_w

mem_size += intermediate_size * ggml_type_sizef(GGML_TYPE_F32); // fc1_b

mem_size += hidden_size * ggml_type_sizef(GGML_TYPE_F32); // fc2_b

mem_size += ggml_type_sizef(GGML_TYPE_F32); // attn_scale

return static_cast<size_t>(mem_size);

}

void init_params(struct ggml_context* ctx, ggml_allocr* alloc, ggml_type wtype) {

ln1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

ln1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

q_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size);

q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

k_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size);

k_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

v_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size);

v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

out_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size);

out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

fc1_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, intermediate_size);

fc1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, intermediate_size);

fc2_w = ggml_new_tensor_2d(ctx, wtype, intermediate_size, hidden_size);

fc2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

ln2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

ln2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size);

attn_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);

ggml_allocr_alloc(alloc, attn_scale);

float scale = 1.0f / sqrt((float)d_model);

ggml_backend_tensor_set(attn_scale, &scale, 0, sizeof(scale));

}

void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {

tensors[prefix + "self_attn.q_proj.weight"] = q_w;

tensors[prefix + "self_attn.q_proj.bias"] = q_b;

tensors[prefix + "self_attn.k_proj.weight"] = k_w;

tensors[prefix + "self_attn.k_proj.bias"] = k_b;

tensors[prefix + "self_attn.v_proj.weight"] = v_w;

tensors[prefix + "self_attn.v_proj.bias"] = v_b;

tensors[prefix + "self_attn.out_proj.weight"] = out_w;

tensors[prefix + "self_attn.out_proj.bias"] = out_b;

tensors[prefix + "layer_norm1.weight"] = ln1_w;

tensors[prefix + "layer_norm1.bias"] = ln1_b;

tensors[prefix + "layer_norm2.weight"] = ln2_w;

tensors[prefix + "layer_norm2.bias"] = ln2_b;

tensors[prefix + "mlp.fc1.weight"] = fc1_w;

tensors[prefix + "mlp.fc1.bias"] = fc1_b;

tensors[prefix + "mlp.fc2.weight"] = fc2_w;

tensors[prefix + "mlp.fc2.bias"] = fc2_b;

}

struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) {

// x: [N, n_token, hidden_size]

int64_t N = x->ne[2];

int64_t n_token = x->ne[1];

int64_t hidden_size = n_head * d_model;

struct ggml_tensor* r = x;

// layer norm 1

x = ggml_nn_layer_norm(ctx, x, ln1_w, ln1_b);

// self-attention

{

struct ggml_tensor* q = ggml_nn_linear(ctx, x, q_w, q_b);

q = ggml_scale_inplace(ctx, q, attn_scale);

q = ggml_reshape_4d(ctx, q, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model]

q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, n_token, d_model]

q = ggml_reshape_3d(ctx, q, d_model, n_token, n_head * N); // [N * n_head, n_token, d_model]

struct ggml_tensor* k = ggml_nn_linear(ctx, x, k_w, k_b);

k = ggml_reshape_4d(ctx, k, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model]

k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, n_token, d_model]

k = ggml_reshape_3d(ctx, k, d_model, n_token, n_head); // [N * n_head, n_token, d_model]

struct ggml_tensor* v = ggml_nn_linear(ctx, x, v_w, v_b);

v = ggml_reshape_4d(ctx, v, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model]

v = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_head, d_model, n_token]

v = ggml_reshape_3d(ctx, v, n_token, d_model, n_head * N); // [N * n_head, d_model, n_token]

struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q); // [N * n_head, n_token, n_token]

kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);

kq = ggml_soft_max_inplace(ctx, kq);

struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, n_token, d_model]

kqv = ggml_reshape_4d(ctx, kqv, d_model, n_token, n_head, N);

kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, n_token, n_head, d_model]