UNPKG

@hoff97/tensor-js

Version:

PyTorch like deep learning inferrence library

49 lines (48 loc) 1.52 kB
import { defaultAllocator } from '../../../tensor/gpu/gl'; import { getSize } from '../../../util/shape'; import { Operation } from '../operation'; export class UnaryOperation extends Operation { constructor(tensorConstructor, dtype, allocator) { super(tensorConstructor, dtype, allocator); } // eslint-disable-next-line @typescript-eslint/no-unused-vars getFragmentShader(info) { return ` void main() { initVars(); gl_FragColor = ${this.operation('texture2D(X, uv)')}; } `; } getTextureNames() { return ['X']; } calc(input) { return this.compute(input.input.shape, { X: input.input }); } getOutputShape(input) { return input.input.shape; } compile(info) { if (info.shapeX !== undefined) { this.maxRank = info.shapeX.length; } super.compile(info); } getCompilationInfo(input) { const outputShape = this.getOutputShape(input); const outputSize = defaultAllocator.getAllocationDimensions(getSize(outputShape), this.dtype); return { shapeX: input.input.shape, widthX: input.input.memory.width, heightX: input.input.memory.height, shapeOutput: this.getOutputShape(input), widthOutput: outputSize.width, heightOutput: outputSize.height, }; } getInputInfoString(input) { return `${input.input.shape}`; } } //# sourceMappingURL=unaryOperation.js.map