ckernel_embedding_checked.h File Reference

#include <stddef.h>
#include <stdint.h>

Go to the source code of this file.

Functions
int	embedding_three_table_layer_norm_f32 (const int32_t *word_ids, const int32_t *type_ids, size_t id_elements, const float *word, size_t word_elements, size_t vocab_size, const float *position, size_t position_elements, size_t position_count, const float *token_type, size_t type_elements, size_t type_count, const float *gamma, const float *beta, size_t norm_elements, float *output, size_t output_elements, size_t tokens, size_t channels, size_t weight_stride, size_t output_stride, float epsilon)

Function Documentation

embedding_three_table_layer_norm_f32()

int embedding_three_table_layer_norm_f32	(	const int32_t *	word_ids,
		const int32_t *	type_ids,
		size_t	id_elements,
		const float *	word,
		size_t	word_elements,
		size_t	vocab_size,
		const float *	position,
		size_t	position_elements,
		size_t	position_count,
		const float *	token_type,
		size_t	type_elements,
		size_t	type_count,
		const float *	gamma,
		const float *	beta,
		size_t	norm_elements,
		float *	output,
		size_t	output_elements,
		size_t	tokens,
		size_t	channels,
		size_t	weight_stride,
		size_t	output_stride,
		float	epsilon
	)

Definition at line 16 of file embedding_three_table_layer_norm.c.

{
    if (!word_ids || !type_ids || !word || !position || !token_type ||
        !gamma || !beta || !output || !tokens || !channels ||
        id_elements < tokens || norm_elements < channels ||
        !isfinite(epsilon) || epsilon <= 0.0f ||
        !span_fits(vocab_size, weight_stride, channels, word_elements) ||
        !span_fits(position_count, weight_stride, channels, position_elements) ||
        !span_fits(type_count, weight_stride, channels, type_elements) ||
        !span_fits(tokens, output_stride, channels, output_elements) ||
        tokens > position_count)
        return -1;
 
    /* Validate every input before the first output write. */
    for (size_t c = 0; c < channels; ++c)
        if (!isfinite(gamma[c]) || !isfinite(beta[c])) return -1;
    for (size_t t = 0; t < tokens; ++t) {
        if (word_ids[t] < 0 || (size_t)word_ids[t] >= vocab_size ||
            type_ids[t] < 0 || (size_t)type_ids[t] >= type_count)
            return -1;
        const size_t wi = (size_t)word_ids[t] * weight_stride;
        const size_t ti = (size_t)type_ids[t] * weight_stride;
        const size_t pi = t * weight_stride;
        for (size_t c = 0; c < channels; ++c)
            if (!isfinite(word[wi + c]) || !isfinite(token_type[ti + c]) ||
                !isfinite(position[pi + c])) return -1;
    }
 
    for (size_t t = 0; t < tokens; ++t) {
        const float *w = word + (size_t)word_ids[t] * weight_stride;
        const float *ty = token_type + (size_t)type_ids[t] * weight_stride;
        const float *p = position + t * weight_stride;
        double sum = 0.0, square_sum = 0.0;
        for (size_t c = 0; c < channels; ++c) {
            const float a = w[c] + ty[c];
            const float x = a + p[c];
            sum += (double)x;
            square_sum += (double)x * (double)x;
        }
        const double mean = sum / (double)channels;
        const double variance = fmax(0.0, square_sum / (double)channels - mean * mean);
        const float inv_std = (float)(1.0 / sqrt(variance + (double)epsilon));
        float *dst = output + t * output_stride;
        for (size_t c = 0; c < channels; ++c) {
            const float a = w[c] + ty[c];
            const float x = a + p[c];
            dst[c] = ((x - (float)mean) * inv_std) * gamma[c] + beta[c];
        }
    }
    return 0;
}

References span_fits(), and vocab_size.