@hoff97/tensor-js/dist/lib/ops/sparse/binary/cpu.js

PyTorch like deep learning inferrence library

import { CPUTensor } from '../../../tensor/cpu/tensor';
import { SparseTensor } from '../../../tensor/sparse/tensor';
import { computeStrides, getSize, incrementIndex } from '../../../util/shape';
export function binaryDenseCPU(a, b, resultShape, op) {
    const S = a.sparseDims;
    const sparseResultShape = resultShape.slice(0, S);
    const denseResultShape = resultShape.slice(S);
    const denseSize = getSize(denseResultShape, 1);
    const valsA = a.values;
    const indicesA = a.indices;
    const values = new CPUTensor([a.nnz, ...denseResultShape], undefined, a.dtype);
    const indices = a.indices.copy();
    for (let i = 0; i < a.nnz; i++) {
        const sparseIx = new Array(sparseResultShape.length);
        for (let j = 0; j < sparseResultShape.length; j++) {
            sparseIx[j] = indicesA.get(i * S + j);
        }
        const denseIx = new Array(denseResultShape.length).fill(0);
        for (let j = 0; j < denseSize; j++) {
            const vA = valsA.get([i, ...denseIx]);
            const vB = b.get([...sparseIx, ...denseIx]);
            values.set(i * denseSize + j, op(vA, vB));
            incrementIndex(denseIx, denseResultShape);
        }
    }
    return new SparseTensor(values, indices, resultShape, a.denseDims);
}
export function binarySparseCPU(a, b, resultShape, op) {
    const S = a.sparseDims;
    const sparseResultShape = resultShape.slice(0, S);
    const denseResultShape = resultShape.slice(S);
    const denseSize = getSize(denseResultShape, 1);
    const valsA = a.values;
    const indicesA = a.indices;
    const valsB = b.values;
    const indicesB = b.indices;
    const sparseStrides = computeStrides(sparseResultShape);
    const bIxPositionMap = {};
    for (let i = 0; i < b.nnz; i++) {
        let pos = 0;
        for (let j = 0; j < S; j++) {
            pos += indicesB.get(i * S + j) * sparseStrides[j];
        }
        bIxPositionMap[pos] = i;
    }
    const values = new CPUTensor([a.nnz, ...denseResultShape], undefined, a.dtype);
    const indices = a.indices.copy();
    for (let i = 0; i < a.nnz; i++) {
        let pos = 0;
        for (let j = 0; j < S; j++) {
            pos += indicesA.get(i * S + j) * sparseStrides[j];
        }
        const iB = bIxPositionMap[pos];
        const denseIx = new Array(denseResultShape.length).fill(0);
        for (let j = 0; j < denseSize; j++) {
            const vA = valsA.get([i, ...denseIx]);
            const vB = valsB.get([iB, ...denseIx]);
            values.set(i * denseSize + j, op(vA, vB));
            incrementIndex(denseIx, denseResultShape);
        }
    }
    return new SparseTensor(values, indices, resultShape, a.denseDims);
}
//# sourceMappingURL=cpu.js.map