onnxruntime/js/web/lib/onnxjs/ops/slice.ts

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

import {Attribute} from '../attribute';
import {InferenceHandler} from '../backend';
import {NUMBER_TYPES, Operator} from '../operators';
import {Tensor} from '../tensor';

export abstract class Slice implements Operator {
  abstract run(inferenceHandler: InferenceHandler, inputs: Tensor[]): Tensor[]|Promise<Tensor[]>;

  initialize(attributes: Attribute): void {
    this.starts = attributes.getInts('starts');
    this.ends = attributes.getInts('ends');
    this.axes = attributes.getInts('axes', []);
  }

  checkInputs(inputs: Tensor[]): boolean {
    if (!inputs || inputs.length !== 1) {
      return false;
    }
    return this.checkInputTypes(inputs);
  }

  protected checkInputTypes(inputs: Tensor[]): boolean {
    if (NUMBER_TYPES.indexOf(inputs[0].type) === -1) {
      return false;
    }
    return true;
  }

  protected axes: number[];
  protected ends: number[];
  protected starts: number[];
}

export abstract class SliceV10 implements Operator {
  abstract run(inferenceHandler: InferenceHandler, inputs: Tensor[]): Tensor[]|Promise<Tensor[]>;

  initialize(_attributes: Attribute): void {}

  checkInputs(inputs: Tensor[]): boolean {
    if (!inputs || inputs.length < 3 || inputs.length > 5) {
      return false;
    }
    return this.checkInputTypes(inputs);
  }

  protected checkInputTypes(inputs: Tensor[]): boolean {
    if (inputs[1].type !== 'int32' || inputs[1].dims.length !== 1) {
      return false;
    }
    if (inputs[2].type !== 'int32' || inputs[2].dims.length !== 1) {
      return false;
    }
    if (inputs.length >= 4 && (inputs[3].type !== 'int32' || inputs[3].dims.length !== 1)) {
      return false;
    }
    if (inputs.length >= 5 && (inputs[4].type !== 'int32' || inputs[4].dims.length !== 1)) {
      return false;
    }

    return true;
  }
}