@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
49 lines (48 loc) • 1.52 kB
JavaScript
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