UNPKG

@hoff97/tensor-js

Version:

PyTorch like deep learning inferrence library

32 lines 1.49 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 { OnnxNode } from '../node'; export class GatherNode extends OnnxNode { constructor(attributes, inputs, outputs, constants, onnxVersion, mode) { super(attributes, inputs, outputs, constants, onnxVersion, mode); this.axis = this.getAttributeInt('axis') || 0; } forward(inputs) { return __awaiter(this, void 0, void 0, function* () { const x = inputs[0]; const indices = inputs[1]; if (!(indices instanceof CPUTensor)) { throw new Error('Gather requires CPU tensor for the indices'); } return [x.gather(this.axis, indices)]; }); } getType() { return 'Gather'; } delete() { } } //# sourceMappingURL=gather.js.map