Function flow_attention

src/models/heartmula/codec.cpp:795–832

        GGML_TYPE_F32);
    context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), core::TensorShape::from_dims({batch, frames, hidden_size}));
    return modules::LinearModule({hidden_size, hidden_size, false, GGML_PREC_F32}).build(ctx, context, weights.o_proj);
}

core::TensorValue flow_mlp(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const HeartCodecTransformerBlockWeights & weights,
    int64_t hidden_size) {
    const int64_t mlp_hidden_size = weights.mlp_gate.weight.shape.dims[0];
    auto gate = modules::LinearModule({hidden_size, mlp_hidden_size, false, GGML_PREC_F32}).build(ctx, input, weights.mlp_gate);
    gate = modules::SiluModule{}.build(ctx, gate);
    auto up = modules::LinearModule({hidden_size, mlp_hidden_size, false, GGML_PREC_F32}).build(ctx, input, weights.mlp_up);
    auto hidden = modules::MulModule{}.build(ctx, gate, up);
    return modules::LinearModule({mlp_hidden_size, hidden_size, false, GGML_PREC_F32}).build(ctx, hidden, weights.mlp_down);
}

core::TensorValue flow_transformer_block(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & timestep_mod,
    const core::TensorValue & positions,
    const HeartCodecTransformerBlockWeights & weights,
    int64_t heads,
    int64_t head_dim) {
    const int64_t hidden_size = heads * head_dim;
    const int64_t frames = input.shape.dims[1];
    const auto parts = adaptive_block_parts(ctx, timestep_mod, weights.scale_shift_table, frames, hidden_size);

    auto normed = modules::RMSNormModule({hidden_size, 1.0e-6F, true, false}).build(ctx, input, weights.attn_norm);
    normed = modules::AddModule{}.build(
        ctx,
        modules::MulModule{}.build(ctx, normed, add_one(ctx, parts.scale_msa)),
        parts.shift_msa);
    auto attn = flow_attention(ctx, normed, positions, weights, heads, head_dim);
    auto x = modules::AddModule{}.build(ctx, input, modules::MulModule{}.build(ctx, parts.gate_msa, attn));

    normed = modules::RMSNormModule({hidden_size, 1.0e-6F, true, false}).build(ctx, x, weights.mlp_norm);
    normed = modules::AddModule{}.build(
        ctx,
        modules::MulModule{}.build(ctx, normed, add_one(ctx, parts.scale_mlp)),
        parts.shift_mlp);