onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { Graph } from '../../../graph';
5
import { OperatorImplementation, OperatorInitialization } from '../../../operators';
6
import { Tensor } from '../../../tensor';
7
import { ShapeUtil } from '../../../util';
8
import { WebGLInferenceHandler } from '../inference-handler';
9

10
export const squeeze: OperatorImplementation<number[]> = (
11
  inferenceHandler: WebGLInferenceHandler,
12
  inputs: Tensor[],
13
  axes: number[],
14
): Tensor[] => {
15
  validateInputs(inputs);
16
  const outputShape = ShapeUtil.squeezeShape(inputs[0].dims, axes);
17
  const output = inferenceHandler.reshapeUnpacked(inputs[0], outputShape);
18
  return [output];
19
};
20

21
export const squeezeV13 = (inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] => {
22
  validateInputsV13(inputs);
23
  return squeeze(inferenceHandler, [inputs[0]], Array.from(inputs[1].integerData));
24
};
25

26
export const parseSqueezeAttributes: OperatorInitialization<number[]> = (node: Graph.Node): number[] =>
27
  node.attributes.getInts('axes');
28

29
const validateInputs = (inputs: Tensor[]): void => {
30
  if (!inputs || inputs.length !== 1) {
31
    throw new Error('Squeeze requires 1 input.');
32
  }
33

34
  if (inputs[0].type === 'string') {
35
    throw new Error('invalid input tensor types.');
36
  }
37
};
38

39
const validateInputsV13 = (inputs: Tensor[]): void => {
40
  if (!inputs || inputs.length !== 2) {
41
    throw new Error('Squeeze requires 2 inputs.');
42
  }
43

44
  if (inputs[1].type !== 'int32') {
45
    throw new Error('Invalid input type.');
46
  }
47
};
48