Function build_time_mask_4d

src/framework/modules/conformer_modules.cpp:48–71

}

core::TensorValue build_time_mask_4d(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & lengths) {
    core::validate_rank_between(input, 4, 4, "time_mask.input");
    core::validate_shape(lengths, core::TensorShape::from_dims({input.shape.dims[0]}), "time_mask.lengths");

    const int64_t batch = input.shape.dims[0];
    const int64_t channels = input.shape.dims[1];
    const int64_t frames = input.shape.dims[2];
    const int64_t features = input.shape.dims[3];

    auto lengths_f32 = core::wrap_tensor(ggml_cast(ctx.ggml, lengths.tensor, GGML_TYPE_F32), lengths.shape, GGML_TYPE_F32);
    auto lengths_4d = core::reshape_tensor(ctx, lengths_f32, core::TensorShape::from_dims({batch, 1, 1, 1}));
    lengths_4d = RepeatModule(RepeatConfig{core::TensorShape::from_dims({batch, 1, frames, 1})}).build(ctx, lengths_4d);

    auto positions = core::wrap_tensor(ggml_arange(ctx.ggml, 0.5f, static_cast<float>(frames) + 0.5f, 1.0f), core::TensorShape::from_dims({frames}), GGML_TYPE_F32);
    positions = core::reshape_tensor(ctx, positions, core::TensorShape::from_dims({1, 1, frames, 1}));
    positions = RepeatModule(RepeatConfig{core::TensorShape::from_dims({batch, 1, frames, 1})}).build(ctx, positions);

    auto diff = core::wrap_tensor(ggml_add(ctx.ggml, lengths_4d.tensor, ggml_scale(ctx.ggml, positions.tensor, -1.0f)), lengths_4d.shape, GGML_TYPE_F32);
    auto mask = core::wrap_tensor(ggml_step(ctx.ggml, diff.tensor), diff.shape, GGML_TYPE_F32);
    return RepeatModule(RepeatConfig{core::TensorShape::from_dims({batch, channels, frames, features})}).build(ctx, mask);
}

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