Function timestep_embedding

src/models/heartmula/codec.cpp:732–744

};

AdaNormParts adaptive_block_parts(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & timestep_mod,
    const core::TensorValue & scale_shift_table,
    int64_t tokens,
    int64_t hidden_size) {
    const int64_t batch = timestep_mod.shape.dims[0];
    auto scale_shift = core::reshape_tensor(ctx, scale_shift_table, core::TensorShape::from_dims({1, 6, hidden_size}));
    scale_shift = modules::RepeatModule({core::TensorShape::from_dims({batch, 6, hidden_size})}).build(ctx, scale_shift);
    auto parts = modules::AddModule{}.build(
        ctx,
        core::reshape_tensor(
            ctx,
            core::ensure_backend_addressable_layout(ctx, timestep_mod),
            core::TensorShape::from_dims({batch, 6, hidden_size})),
        scale_shift);