Function global_encoder

src/models/miocodec/graph_ops.cpp:586–621

}

core::TensorValue global_encoder(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input_btc,
    const MioCodecGlobalEncoderWeights & weights) {
    auto x = engine::modules::TransposeModule({{0, 2, 1, 3}, input_btc.shape.rank}).build(ctx, input_btc);
    x = engine::modules::Conv1dModule(weights.embedding.config).build(ctx, x, weights.embedding.weights);
    x = engine::modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(ctx, x);
    x = engine::modules::LayerNormModule({384, 1.0e-6F, true, true}).build(ctx, x, weights.embedding_norm);
    x = engine::modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(ctx, x);
    for (const auto & block : weights.blocks) {
        x = convnext_block(ctx, x, block);
    }
    auto features_btc = engine::modules::TransposeModule({{0, 2, 1, 3}, x.shape.rank}).build(ctx, x);
    features_btc = engine::modules::LayerNormModule({384, 1.0e-6F, true, true}).build(ctx, features_btc, weights.final_norm);
    auto features_bct = engine::modules::TransposeModule({{0, 2, 1, 3}, features_btc.shape.rank}).build(ctx, features_btc);
    auto alpha = engine::modules::Conv1dModule(weights.attention_conv1.config).build(ctx, features_bct, weights.attention_conv1.weights);
    alpha = core::wrap_tensor(ggml_tanh(ctx.ggml, alpha.tensor), alpha.shape, GGML_TYPE_F32);
    alpha = engine::modules::Conv1dModule(weights.attention_conv2.config).build(ctx, alpha, weights.attention_conv2.weights);
    alpha = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, alpha).tensor), alpha.shape, GGML_TYPE_F32);
    auto weighted = core::wrap_tensor(ggml_mul(ctx.ggml, alpha.tensor, features_bct.tensor), features_bct.shape, GGML_TYPE_F32);
    auto mean = engine::modules::ReduceSumModule({2}).build(ctx, weighted);
    auto features_contiguous = core::ensure_backend_addressable_layout(ctx, features_bct);
    auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, features_contiguous.tensor), features_contiguous.shape, GGML_TYPE_F32);
    auto second = engine::modules::ReduceSumModule({2}).build(ctx, core::wrap_tensor(ggml_mul(ctx.ggml, alpha.tensor, squared.tensor), squared.shape, GGML_TYPE_F32));
    auto mean_sq = core::wrap_tensor(ggml_sqr(ctx.ggml, mean.tensor), mean.shape, GGML_TYPE_F32);
    auto var = core::wrap_tensor(ggml_sub(ctx.ggml, second.tensor, mean_sq.tensor), mean.shape, GGML_TYPE_F32);
    var = core::wrap_tensor(ggml_clamp(ctx.ggml, var.tensor, 1.0e-4F, 1.0e4F), var.shape, GGML_TYPE_F32);
    auto std = core::wrap_tensor(ggml_sqrt(ctx.ggml, var.tensor), var.shape, GGML_TYPE_F32);
    auto stats = engine::modules::ConcatModule({1}).build(ctx, mean, std);
    stats = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, stats), core::TensorShape::from_dims({stats.shape.dims[0], 768}));
    auto out = engine::modules::LinearModule({768, 128, weights.pooling_projection.bias.has_value()}).build(
        ctx,
        stats,
        weights.pooling_projection);
    return engine::modules::LayerNormModule({128, 1.0e-5F, true, true}).build(ctx, out, weights.pooling_norm);
}

core::TensorValue dynamic_linear_interpolate_bct(
    core::ModuleBuildContext & ctx,