core::ModuleBuildContext & ctx,

const core::TensorValue & value) {

return core::ensure_backend_addressable_layout(ctx, value);

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const ZipformerLinearWeights & weights,

const char * name) {

core::validate_rank_between(weights.weight, 2, 2, name);

if (weights.weight.shape.dims[1] != input.shape.last_dim()) {

throw std::runtime_error(

std::string(name) + " input dimension mismatch");

}

return LinearModule({

weights.weight.shape.dims[1],

weights.weight.shape.dims[0],

true,

}).build(ctx, input, {weights.weight, weights.bias});

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const std::array<int, core::kMaxTensorRank> & axes) {

const size_t rank = input.shape.rank;

core::TensorShape output = {};

output.rank = rank;

std::array<bool, core::kMaxTensorRank> seen = {false, false, false, false};

std::array<int, core::kMaxTensorRank> ggml_axes = {0, 1, 2, 3};

for (size_t out_axis = 0; out_axis < rank; ++out_axis) {

const int in_axis = axes[out_axis];

if (in_axis < 0 ||

in_axis >= static_cast<int>(rank) ||

seen[static_cast<size_t>(in_axis)]) {

throw std::runtime_error("invalid Zipformer transpose");

}

seen[static_cast<size_t>(in_axis)] = true;

output.dims[out_axis] = input.shape.dims[in_axis];

const int out_ggml =

core::logical_axis_to_ggml_axis(rank, out_axis);

const int in_ggml =

core::logical_axis_to_ggml_axis(rank, in_axis);

ggml_axes[static_cast<size_t>(in_ggml)] = out_ggml;

}

const auto source = contiguous(ctx, input);

return core::wrap_tensor(

ggml_permute(

ctx.ggml,

source.tensor,

ggml_axes[0],

ggml_axes[1],

ggml_axes[2],

ggml_axes[3]),

output,

input.type);

int64_t seq,

int64_t left_context,

int64_t position_dim) {

constexpr float pi = 3.14159265358979323846F;

const int64_t length = left_context + 2 * seq - 1;

std::vector<float> result(

static_cast<size_t>(length * position_dim), 0.0F);

const float compression =

std::sqrt(static_cast<float>(position_dim));

const float length_scale =

static_cast<float>(position_dim) / (2.0F * pi);

for (int64_t row = 0; row < length; ++row) {

const float x =

static_cast<float>(row - (left_context + seq - 1));

const float sign =

x < 0.0F ? -1.0F : (x > 0.0F ? 1.0F : 0.0F);

const float compressed =

compression * sign *

(std::log(std::fabs(x) + compression) -

std::log(compression));

const float angle = std::atan(compressed / length_scale);

for (int64_t dim = 0; dim < position_dim / 2; ++dim) {

const float frequency = static_cast<float>(dim + 1);

result[static_cast<size_t>(row * position_dim + 2 * dim)] =

std::cos(angle * frequency);

result[static_cast<size_t>(row * position_dim + 2 * dim + 1)] =

std::sin(angle * frequency);

}

result[static_cast<size_t>(row * position_dim + position_dim - 1)] =

1.0F;

}

return result;

ZipformerRelativeAttentionConfig config)

: config_(config) {

if (config_.num_heads <= 0 ||

config_.query_head_dim <= 0 ||

config_.position_dim <= 0 ||

config_.position_head_dim <= 0) {

throw std::runtime_error(

"ZipformerRelativeAttentionModule config dimensions must be positive");

}

core::ModuleBuildContext & ctx,

ZipformerConstantFactory & constant_factory,

const core::TensorValue & input,

const core::TensorValue & cached_key,

const core::TensorValue & padding_mask,

const ZipformerRelativeAttentionWeights & weights) const {

const int64_t seq = input.shape.dims[0];

const int64_t left_context = cached_key.shape.dims[0];

const int64_t source_frames = left_context + seq;

const int64_t query_channels =

config_.num_heads * config_.query_head_dim;

const int64_t position_channels =

config_.num_heads * config_.position_head_dim;

const size_t expected_position_values =

static_cast<size_t>(

position_channels * config_.position_dim);

if (weights.linear_pos_weight.size() != expected_position_values) {

throw std::runtime_error(

"Zipformer relative position weight size mismatch");

}

auto projected =

linear(ctx, input, weights.in_proj, "Zipformer relative attention in_proj");

auto current_key =

SliceModule({2, query_channels, query_channels})

.build(ctx, projected);

auto all_keys = ConcatModule({0}).build(ctx, cached_key, current_key);

ZipformerRelativeAttentionOutputs result;

result.key_state =

SliceModule({0, seq, left_context}).build(ctx, all_keys);

const auto position =

compact_relative_position(

seq,

left_context,

config_.position_dim);

result.weights.reserve(static_cast<size_t>(config_.num_heads));

for (int64_t head = 0; head < config_.num_heads; ++head) {

auto query =

SliceModule({

2,

head * config_.query_head_dim,

config_.query_head_dim,

}).build(ctx, projected);

auto key =

SliceModule({

2,

head * config_.query_head_dim,

config_.query_head_dim,

}).build(ctx, all_keys);

auto position_query =

SliceModule({

2,

2 * query_channels + head * config_.position_head_dim,

config_.position_head_dim,

}).build(ctx, projected);

query = transpose(ctx, query, {1, 0, 2, 3});

key = transpose(ctx, key, {1, 2, 0, 3});

auto scores = MatMulModule().build(ctx, query, key);

position_query = transpose(ctx, position_query, {1, 0, 2, 3});

core::TensorValue position_scores;

for (int64_t dim = 0; dim < config_.position_head_dim; ++dim) {

std::vector<float> table(

static_cast<size_t>(seq * source_frames), 0.0F);

const int64_t output_row =

head * config_.position_head_dim + dim;

for (int64_t target = 0; target < seq; ++target) {

for (int64_t source = 0;

source < source_frames;

++source) {

const int64_t relative =

(seq - 1 - target) + source;

double value = 0.0;

for (int64_t p = 0; p < config_.position_dim; ++p) {

value +=

static_cast<double>(

position[static_cast<size_t>(

relative * config_.position_dim + p)]) *

static_cast<double>(

weights.linear_pos_weight[

static_cast<size_t>(

output_row * config_.position_dim + p)]);

}

table[static_cast<size_t>(

target * source_frames + source)] =

static_cast<float>(value);

}

}

auto relative = constant_factory(

core::TensorShape::from_dims({1, seq, source_frames}),

std::move(table));

auto query_dim_value =

SliceModule({2, dim, 1}).build(ctx, position_query);

auto repeated =

RepeatModule({

core::TensorShape::from_dims(

{1, seq, source_frames}),

}).build(ctx, query_dim_value);

auto term = MulModule().build(ctx, repeated, relative);

position_scores = position_scores.valid()

? AddModule().build(ctx, position_scores, term)

: term;

}

result.weights.push_back(

SoftmaxModule().build(

ctx,

AddModule().build(

ctx,

AddModule().build(ctx, scores, position_scores),

padding_mask)));

}

return result;

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const core::TensorValue & attention,

const core::TensorValue & cached_value,

const ZipformerNonlinearAttentionWeights & weights) const {

auto projected =

linear(ctx, input, weights.in_proj, "Zipformer nonlinear attention in_proj");

const int64_t hidden = projected.shape.last_dim() / 3;

auto gate_input =

SliceModule({2, 0, hidden}).build(ctx, projected);

auto value =

SliceModule({2, hidden, hidden}).build(ctx, projected);

auto output_gate =

SliceModule({2, 2 * hidden, hidden}).build(ctx, projected);

value = MulModule().build(

ctx,

value,

TanhModule().build(ctx, gate_input));

const int64_t seq = value.shape.dims[0];

const int64_t left_context = cached_value.shape.dims[0];

value = ConcatModule({0}).build(ctx, cached_value, value);

auto new_state =

SliceModule({0, seq, left_context}).build(ctx, value);

auto joined = transpose(ctx, value, {1, 0, 2, 3});

joined = MatMulModule().build(ctx, attention, joined);

joined = transpose(ctx, joined, {1, 0, 2, 3});

return {

linear(

ctx,

MulModule().build(ctx, joined, output_gate),

weights.out_proj,

"Zipformer nonlinear attention out_proj"),

new_state,

};

ZipformerConvolutionConfig config)

: config_(config) {

if (config_.kernel_size <= 0) {

throw std::runtime_error(

"ZipformerConvolutionModule kernel size must be positive");

}

core::ModuleBuildContext & ctx,

ZipformerConstantFactory & constant_factory,

const core::TensorValue & input,

const core::TensorValue & cached_value,

const ZipformerConvolutionWeights & weights) const {

const int64_t channels = input.shape.last_dim();

auto projected =

linear(ctx, input, weights.in_proj, "Zipformer convolution in_proj");

auto value =

SliceModule({2, 0, channels}).build(ctx, projected);

auto gate =

SliceModule({2, channels, channels}).build(ctx, projected);

value = MulModule().build(

ctx,

value,

SigmoidModule().build(ctx, gate));

auto channel_first = transpose(ctx, value, {1, 2, 0, 3});

const int64_t left_pad = config_.kernel_size / 2;

auto causal_input =

ConcatModule({2}).build(ctx, cached_value, channel_first);

auto new_state =

contiguous(

ctx,

SliceModule({2, input.shape.dims[0], left_pad})

.build(ctx, causal_input));

auto causal =

DepthwiseConv1dModule(

{channels, left_pad + 1, 1, 0, 1, true})

.build(ctx, causal_input, weights.causal_conv);

auto chunk =

DepthwiseConv1dModule({

channels,

config_.kernel_size,

1,

static_cast<int>(left_pad),

1,

true,

}).build(ctx, channel_first, weights.chunkwise_conv);

if (weights.chunk_scale_left.size() !=

static_cast<size_t>(channels * config_.kernel_size) ||

weights.chunk_scale_right.size() !=

static_cast<size_t>(channels * config_.kernel_size)) {

throw std::runtime_error(

"Zipformer convolution chunk scale size mismatch");

}

const int64_t seq = input.shape.dims[0];

std::vector<float> chunk_scale(

static_cast<size_t>(channels * seq), 1.0F);

for (int64_t channel = 0; channel < channels; ++channel) {

for (int64_t frame = 0; frame < seq; ++frame) {

float edge = 0.0F;

if (seq < config_.kernel_size) {

edge += weights.chunk_scale_left[static_cast<size_t>(

channel * config_.kernel_size + frame)];

edge += weights.chunk_scale_right[static_cast<size_t>(

channel * config_.kernel_size +

config_.kernel_size - seq + frame)];

} else {

if (frame < config_.kernel_size) {

edge += weights.chunk_scale_left[static_cast<size_t>(

channel * config_.kernel_size + frame)];

}

if (frame >= seq - config_.kernel_size) {

edge += weights.chunk_scale_right[static_cast<size_t>(

channel * config_.kernel_size +

frame - (seq - config_.kernel_size))];

}

}

chunk_scale[static_cast<size_t>(channel * seq + frame)] += edge;

}

}

auto scale_tensor = constant_factory(

core::TensorShape::from_dims({1, channels, seq}),

std::move(chunk_scale));

auto scaled_chunk = MulModule().build(ctx, chunk, scale_tensor);

auto combined = AddModule().build(ctx, causal, scaled_chunk);

combined = transpose(ctx, combined, {2, 0, 1, 3});

combined = SwooshRModule().build(ctx, combined);

return {

linear(

ctx,

combined,

weights.out_proj,

"Zipformer convolution out_proj"),

new_state,

};