Function layer_norm

packages/transformers/src/utils/tensor.js:1205–1233

(input, normalized_shape, { eps = 1e-5 } = {})

 * @returns {Tensor} The normalized tensor.
 */
export function layer_norm(input, normalized_shape, { eps = 1e-5 } = {}) {
    if (input.dims.length !== 2) {
        throw new Error('`layer_norm` currently only supports 2D input.');
    }

    const [batchSize, featureDim] = input.dims;

    if (normalized_shape.length !== 1 && normalized_shape[0] !== featureDim) {
        throw new Error('`normalized_shape` must be a 1D array with shape `[input.dims[1]]`.');
    }

    const [std, mean] = std_mean(input, 1, 0, true);
    const stdData = /** @type {Float32Array} */ (std.data);
    const meanData = /** @type {Float32Array} */ (mean.data);

    const inputData = /** @type {Float32Array} */ (input.data);

    // @ts-ignore
    const returnedData = new inputData.constructor(inputData.length);

    for (let i = 0; i < batchSize; ++i) {
        const offset = i * featureDim;
        for (let j = 0; j < featureDim; ++j) {
            const offset2 = offset + j;
            returnedData[offset2] = (inputData[offset2] - meanData[i]) / (stdData[i] + eps);
        }
    }
    return new Tensor(input.type, returnedData, input.dims);
}

/**
 * Helper function to calculate new dimensions when performing a squeeze operation.