@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
86 lines (83 loc) • 2.78 kB
JavaScript
import { defaultAllocator } from '../../../tensor/gpu/gl';
import { Operation } from '../operation';
export class ClipOperation 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();
vec4 maxVec = vec4(maxVal,maxVal,maxVal,maxVal);
vec4 minVec = vec4(minVal,minVal,minVal,minVal);
vec4 res = texture2D(X, uv);
if (doMin == 1) {
res = max(minVec, res);
}
if (doMax == 1) {
res = min(maxVec, res);
}
gl_FragColor = res;
}
`;
}
getTextureNames() {
return ['X'];
}
getVariables() {
return `
${this.getVarModifier('minVal')} float minVal;
${this.getVarModifier('maxVal')} float maxVal;
${this.getVarModifier('doMin')} int doMin;
${this.getVarModifier('doMax')} int doMax;
`;
}
getUniformAttrs() {
return [
{ name: 'minVal', type: 'float' },
{ name: 'maxVal', type: 'float' },
{ name: 'doMin' },
{ name: 'doMax' },
];
}
calc(input) {
if (this.fullyStatic && this.outputShape !== undefined) {
return this.compute(this.outputShape, { X: input.input });
}
return this.compute(input.input.shape, { X: input.input }, {
minVal: input.minVal !== undefined ? input.minVal : 0,
maxVal: input.maxVal !== undefined ? input.maxVal : 0,
doMin: input.minVal !== undefined ? 1 : 0,
doMax: input.maxVal !== undefined ? 1 : 0,
});
}
getOutputShape(input) {
return input.input.shape;
}
compile(info) {
if (info.shapeX !== undefined) {
this.maxRank = info.shapeX.length;
}
super.compile(info);
}
getCompilationInfo(input) {
const outputSize = defaultAllocator.getAllocationDimensions(input.input.size, this.dtype);
return {
shapeX: input.input.shape,
widthX: input.input.memory.width,
heightX: input.input.memory.height,
shapeOutput: input.input.shape,
widthOutput: outputSize.width,
heightOutput: outputSize.height,
minVal: input.minVal !== undefined ? input.minVal : 0,
maxVal: input.maxVal !== undefined ? input.maxVal : 0,
doMin: input.minVal !== undefined ? 1 : 0,
doMax: input.maxVal !== undefined ? 1 : 0,
};
}
getInputInfoString(input) {
return `${input.input.shape}-${input.minVal}-${input.maxVal}`;
}
}
//# sourceMappingURL=clip.js.map