Function std_mean

packages/transformers/src/utils/tensor.js:1438–1485

(input, dim = null, correction = 1, keepdim = false)

 * @returns {Tensor[]} A tuple of (std, mean) tensors.
 */
export function std_mean(input, dim = null, correction = 1, keepdim = false) {
    const inputData = /** @type {Float32Array} */ (input.data);
    const inputDims = input.dims;

    if (dim === null) {
        // None to reduce over all dimensions.
        const sum = inputData.reduce((a, b) => a + b, 0);
        const mean = sum / inputData.length;
        const std = Math.sqrt(inputData.reduce((a, b) => a + (b - mean) ** 2, 0) / (inputData.length - correction));

        const meanTensor = new Tensor(
            input.type,
            [mean],
            [
                /* scalar */
            ],
        );
        const stdTensor = new Tensor(
            input.type,
            [std],
            [
                /* scalar */
            ],
        );

        return [stdTensor, meanTensor];
    }
    dim = safeIndex(dim, inputDims.length);
    const meanTensor = mean(input, dim, keepdim);
    const meanTensorData = meanTensor.data;

    // Compute squared sum
    const [type, result, resultDims] = reduce_helper(
        (a, b, i, j) => a + (b - meanTensorData[j]) ** 2,
        input,
        dim,
        keepdim,
    );

    // Square root of the squared sum
    for (let i = 0; i < result.length; ++i) {
        result[i] = Math.sqrt(result[i] / (inputDims[dim] - correction));
    }

    const stdTensor = new Tensor(type, result, resultDims);

    return [stdTensor, meanTensor];
}

/**
 * Returns the mean value of each row of the input tensor in the given dimension dim.