slice()

in src/utils/tensor.js [387:456]


    slice(...slices) {
        // This allows for slicing with ranges and numbers
        const newTensorDims = [];
        const newOffsets = [];

        // slices is an array of numbers or arrays of numbers
        // e.g., slices = [0, [1, 3], null, [0, 3]]
        for (let sliceIndex = 0; sliceIndex < this.dims.length; ++sliceIndex) {
            let slice = slices[sliceIndex];

            if (slice === null || slice === undefined) {
                // null or undefined means take the whole dimension
                newOffsets.push([0, this.dims[sliceIndex]]);
                newTensorDims.push(this.dims[sliceIndex]);

            } else if (typeof slice === 'number') {
                slice = safeIndex(slice, this.dims[sliceIndex], sliceIndex);

                // A number means take a single element
                newOffsets.push([slice, slice + 1]);

            } else if (Array.isArray(slice) && slice.length === 2) {
                // An array of length 2 means take a range of elements
                let [start, end] = slice;
                start = start === null
                    ? 0
                    : safeIndex(start, this.dims[sliceIndex], sliceIndex, false);
                end = end === null
                    ? this.dims[sliceIndex]
                    : safeIndex(end, this.dims[sliceIndex], sliceIndex, false);

                if (start > end) {
                    throw new Error(`Invalid slice: ${slice}`);
                }

                const offsets = [
                    Math.max(start, 0),
                    Math.min(end, this.dims[sliceIndex])
                ];

                newOffsets.push(offsets);
                newTensorDims.push(offsets[1] - offsets[0]);

            } else {
                throw new Error(`Invalid slice: ${slice}`);
            }
        }

        const newDims = newOffsets.map(([start, end]) => end - start);
        const newBufferSize = newDims.reduce((a, b) => a * b);

        const this_data = this.data;
        // Allocate memory
        // @ts-ignore
        const data = new this_data.constructor(newBufferSize);

        // Precompute strides
        const stride = this.stride();

        for (let i = 0; i < newBufferSize; ++i) {
            let originalIndex = 0;
            for (let j = newDims.length - 1, num = i; j >= 0; --j) {
                const size = newDims[j];
                originalIndex += ((num % size) + newOffsets[j][0]) * stride[j];
                num = Math.floor(num / size);
            }
            data[i] = this_data[originalIndex];
        }
        return new Tensor(this.type, data, newTensorDims);
    }