UNPKG

@hoff97/tensor-js

Version:

PyTorch like deep learning inferrence library

53 lines 2.48 kB
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) { function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); } return new (P || (P = Promise))(function (resolve, reject) { function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } } function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } } function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); } step((generator = generator.apply(thisArg, _arguments || [])).next()); }); }; import { CPUTensor } from '../../tensor/cpu/tensor'; import { getSize } from '../../util/shape'; import { OnnxNode } from '../node'; import { createTensor } from '../util'; // This does not support gradients right now, mainly because // the forward pass needs to directly access the constant value export class ConstantOfShapeNode extends OnnxNode { constructor(attributes, inputs, outputs, constants, onnxVersion, mode) { super(attributes, inputs, outputs, constants, onnxVersion, mode); if (onnxVersion < 11) { const tensor = this.getAttributeTensor('value'); if (tensor !== null && tensor !== undefined) { this.tensor = createTensor(tensor); } } } forward(inputs) { return __awaiter(this, void 0, void 0, function* () { const _shape = inputs[0]; if (this.onnxVersion < 11 && this.tensor !== undefined) { if (!(_shape instanceof CPUTensor)) { throw new Error('ConstantOfShape needs cpu tensor as shape tensor'); } const shape = new Array(_shape.size); for (let i = 0; i < _shape.size; i++) { shape[i] = _shape.get(i); } const size = getSize(shape); const values = new Float32Array(size).fill(this.tensor.get(0)); return [new CPUTensor(shape, values, this.tensor.dtype)]; } throw new Error(`ConstantOfShape not implemented for onnx version ${this.onnxVersion}`); }); } getType() { return 'ConstantOfShape'; } delete() { if (this.tensor !== undefined) { this.tensor.delete(); } } } //# sourceMappingURL=constantOfShape.js.map