@hoff97/tensor-js/dist/lib/model/functional/bce/bce.js

PyTorch like deep learning inferrence library

import { Variable } from '../../../autograd';
import { CPUTensor } from '../../../tensor/cpu/tensor';
import { GPUTensor } from '../../../tensor/gpu/tensor';
import { WASMTensor } from '../../../tensor/wasm/tensor';
import { sameType } from '../../../util/convert';
import { BCEBack } from './back/back';
import { bce as bceCPU } from './cpu';
import { defaultBCED } from './gpu';
/**
 * Calculates the binary cross entropy loss, given probabilities x
 * and ground truth y. Returns a tensor of the same shape as
 * x. To use for a loss, you have to sum over the result:
 * ```typescript
 * const loss = bce(x,y).sum();
 * ```
 *
 * @param x Probabilities in [0,1]
 * @param y Ground truth labels of the same shape as x.
 */
export function bce(x, y) {
    if (!sameType(x, y)) {
        throw new Error('BCE can only be computed for tensors of the same type');
    }
    if (x instanceof CPUTensor && y instanceof CPUTensor) {
        return bceCPU(x, y);
    }
    else if (x instanceof WASMTensor && y instanceof WASMTensor) {
        return new WASMTensor(x.wasmTensor.bce(y.wasmTensor));
    }
    else if (x instanceof GPUTensor && y instanceof GPUTensor) {
        return defaultBCED.calc({
            A: x,
            B: y,
            outputShape: x.getShape(),
        }, x.dtype);
    }
    else {
        return new Variable(bce(x.value, y.value), {
            noGrad: x.noGrad,
            backEdge: x.noGrad
                ? undefined
                : new BCEBack(x, y),
        });
    }
}
//# sourceMappingURL=bce.js.map