Function mean

packages/transformers/src/utils/tensor.js:1494–1522

(input, dim = null, keepdim = false)

 * @returns {Tensor} A new tensor with means taken along the specified dimension.
 */
export function mean(input, dim = null, keepdim = false) {
    const inputDims = input.dims;
    const inputData = /** @type {Float32Array} */ (input.data);

    if (dim === null) {
        // None to reduce over all dimensions.
        const val = inputData.reduce((a, b) => a + b, 0);
        return new Tensor(
            input.type,
            [val / inputData.length],
            [
                /* scalar */
            ],
        );
    }
    dim = safeIndex(dim, inputDims.length);

    // Compute sum
    const [type, result, resultDims] = reduce_helper((a, b) => a + b, input, dim, keepdim);

    // Divide by number of elements in the dimension
    if (inputDims[dim] !== 1) {
        for (let i = 0; i < result.length; ++i) {
            result[i] /= inputDims[dim];
        }
    }

    return new Tensor(type, result, resultDims);
}

function dimsToStride(dims) {
    const stride = new Array(dims.length);