// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// https://www.w3.org/TR/webnn/

typedef record<USVString, MLOperand> MLNamedOperands;

enum MLInputOperandLayout { "nchw", "nhwc" };

enum MLConv2dFilterOperandLayout { "oihw", "hwio", "ohwi", "ihwo" };

enum MLConvTranspose2dFilterOperandLayout { "iohw", "hwoi", "ohwi" };

enum MLRecurrentNetworkActivation { "relu", "sigmoid", "tanh" };

enum MLRecurrentNetworkDirection { "forward", "backward", "both" };

enum MLLstmWeightLayout { "iofg", "ifgo" };

enum MLGruWeightLayout { "zrn", "rzn" };

dictionary MLOperatorOptions {
  USVString label = "";
};

dictionary MLArgMinMaxOptions : MLOperatorOptions {
  boolean keepDimensions = false;
  // See spec issue https://github.com/webmachinelearning/webnn/issues/653.
  MLOperandDataType outputDataType = "int32";
};

// A spec file was issued for WG discussion:
// https://github.com/webmachinelearning/webnn/issues/481.
// TODO(crbug.com/1502361): Revisit whether the scale and bias operands
// should be required inputs based on WG's consensus.
dictionary MLBatchNormalizationOptions : MLOperatorOptions {
  MLOperand scale;
  MLOperand bias;
  [EnforceRange] unsigned long axis = 1;
  float epsilon = 1e-5;
};

dictionary MLConv2dOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> padding;
  sequence<[EnforceRange] unsigned long> strides;
  sequence<[EnforceRange] unsigned long> dilations;
  [EnforceRange] unsigned long groups = 1;
  MLInputOperandLayout inputLayout = "nchw";
  MLConv2dFilterOperandLayout filterLayout = "oihw";
  MLOperand bias;
};

dictionary MLConvTranspose2dOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> padding;
  sequence<[EnforceRange] unsigned long> strides;
  sequence<[EnforceRange] unsigned long> dilations;
  sequence<[EnforceRange] unsigned long> outputPadding;
  sequence<[EnforceRange] unsigned long> outputSizes;
  [EnforceRange] unsigned long groups = 1;
  MLInputOperandLayout inputLayout = "nchw";
  MLConvTranspose2dFilterOperandLayout filterLayout = "iohw";
  MLOperand bias;
};

dictionary MLCumulativeSumOptions : MLOperatorOptions {
    boolean exclusive = false;
    boolean reversed = false;
};

dictionary MLGatherOptions : MLOperatorOptions {
  [EnforceRange] unsigned long axis = 0;
};

dictionary MLGemmOptions : MLOperatorOptions {
  MLOperand c;
  float alpha = 1.0;
  float beta = 1.0;
  boolean aTranspose = false;
  boolean bTranspose = false;
};

dictionary MLGruOptions : MLOperatorOptions {
  MLOperand bias;
  MLOperand recurrentBias;
  MLOperand initialHiddenState;
  boolean resetAfter = true;
  boolean returnSequence = false;
  MLRecurrentNetworkDirection direction = "forward";
  MLGruWeightLayout layout = "zrn";
  sequence<MLRecurrentNetworkActivation> activations;
};

dictionary MLGruCellOptions : MLOperatorOptions {
  MLOperand bias;
  MLOperand recurrentBias;
  boolean resetAfter = true;
  MLGruWeightLayout layout = "zrn";
  sequence<MLRecurrentNetworkActivation> activations;
};

dictionary MLHardSigmoidOptions : MLOperatorOptions {
  float alpha = 0.2;
  float beta = 0.5;
};

dictionary MLLayerNormalizationOptions : MLOperatorOptions {
  MLOperand scale;
  MLOperand bias;
  sequence<[EnforceRange] unsigned long> axes;
  float epsilon = 1e-5;
};

dictionary MLLeakyReluOptions : MLOperatorOptions {
  float alpha = 0.01;
};

dictionary MLLinearOptions : MLOperatorOptions {
  float alpha = 1.0;
  float beta = 0;
};

dictionary MLLstmOptions : MLOperatorOptions {
  MLOperand bias;
  MLOperand recurrentBias;
  MLOperand peepholeWeight;
  MLOperand initialHiddenState;
  MLOperand initialCellState;
  boolean returnSequence = false;
  MLRecurrentNetworkDirection direction = "forward";
  MLLstmWeightLayout layout = "iofg";
  sequence<MLRecurrentNetworkActivation> activations;
};

dictionary MLLstmCellOptions : MLOperatorOptions {
  MLOperand bias;
  MLOperand recurrentBias;
  MLOperand peepholeWeight;
  MLLstmWeightLayout layout = "iofg";
  sequence<MLRecurrentNetworkActivation> activations;
};

enum MLPaddingMode {
  "constant",
  "edge",
  "reflection"
};

dictionary MLPadOptions : MLOperatorOptions {
  MLPaddingMode mode = "constant";
  float value = 0;
};

enum MLRoundingType {
  "floor",
  "ceil"
};

dictionary MLPool2dOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> windowDimensions;
  sequence<[EnforceRange] unsigned long> padding;
  sequence<[EnforceRange] unsigned long> strides;
  sequence<[EnforceRange] unsigned long> dilations;
  MLInputOperandLayout layout = "nchw";
  MLRoundingType roundingType = "floor";
  sequence<[EnforceRange] unsigned long> outputSizes;
};

dictionary MLClampOptions : MLOperatorOptions {
  float minValue;
  float maxValue;
};

dictionary MLEluOptions : MLOperatorOptions {
  float alpha = 1;
};

dictionary MLInstanceNormalizationOptions : MLOperatorOptions {
  MLOperand scale;
  MLOperand bias;
  float epsilon = 1e-5;
  MLInputOperandLayout layout = "nchw";
};

dictionary MLReduceOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> axes;
  boolean keepDimensions = false;
};

enum MLInterpolationMode {"nearest-neighbor", "linear" };

dictionary MLResample2dOptions : MLOperatorOptions {
  MLInterpolationMode mode = "nearest-neighbor";
  sequence<float> scales;
  sequence<[EnforceRange] unsigned long> sizes;
  sequence<[EnforceRange] unsigned long> axes;
};

dictionary MLReverseOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> axes;
};

dictionary MLScatterOptions : MLOperatorOptions {
  [EnforceRange] unsigned long axis = 0;
};

dictionary MLSliceOptions : MLOperatorOptions {
  [EnforceRange] sequence<[EnforceRange] unsigned long> strides;
};

dictionary MLTransposeOptions : MLOperatorOptions {
  sequence<[EnforceRange] unsigned long> permutation;
};

dictionary MLSplitOptions : MLOperatorOptions {
  [EnforceRange] unsigned long axis = 0;
};

dictionary MLTriangularOptions : MLOperatorOptions {
  boolean upper = true;
  [EnforceRange] long diagonal = 0;
};

[
  RuntimeEnabled=MachineLearningNeuralNetwork,
  Exposed=(Window, Worker)
] interface MLGraphBuilder {
  [
    CallWith=ScriptState,
    RaisesException
  ] constructor(MLContext context);

  [
    CallWith=ScriptState,
    RaisesException
  ] MLOperand input(USVString name, MLOperandDescriptor desc);

  [
    CallWith=ScriptState,
    RaisesException
  ] MLOperand constant(MLOperandDescriptor desc, AllowSharedBufferSource buffer);

  [
    CallWith=ScriptState,
    RaisesException
  ] MLOperand constant(MLTensor tensor);

  [RaisesException] MLOperand argMin(MLOperand input, [EnforceRange] unsigned long axis, optional MLArgMinMaxOptions options = {});
  [RaisesException] MLOperand argMax(MLOperand input, [EnforceRange] unsigned long axis, optional MLArgMinMaxOptions options = {});

  [RaisesException] MLOperand batchNormalization(MLOperand input, MLOperand mean, MLOperand variance, optional MLBatchNormalizationOptions options = {});

  [RaisesException] MLOperand clamp(MLOperand input, optional MLClampOptions options = {});

  [RaisesException] MLOperand concat(sequence<MLOperand> inputs, [EnforceRange] unsigned long axis, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand conv2d(MLOperand input, MLOperand filter, optional MLConv2dOptions options = {});
  [RaisesException] MLOperand convTranspose2d(MLOperand input, MLOperand filter, optional MLConvTranspose2dOptions options = {});

  [RaisesException] MLOperand cumulativeSum(MLOperand input, [EnforceRange] unsigned long axis, optional MLCumulativeSumOptions options = {});

  // Element-wise binary operations
  [RaisesException] MLOperand add(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand sub(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand mul(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand div(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand max(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand min(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand pow(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand equal(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand greater(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand greaterOrEqual(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand lesser(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand lesserOrEqual(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand notEqual(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand logicalAnd(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand logicalOr(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand logicalXor(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});

  // Element-wise unary operations
  [RaisesException] MLOperand abs(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand ceil(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand cos(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand exp(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand floor(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand log(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand neg(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand sign(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand sin(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand tan(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand erf(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand identity(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand logicalNot(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand reciprocal(MLOperand x, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand sqrt(MLOperand x, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand cast(MLOperand input, MLOperandDataType outputDataType, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand dequantizeLinear(MLOperand input, MLOperand scale, MLOperand zeroPoint, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand elu(MLOperand x, optional MLEluOptions options = {});

  [RaisesException] MLOperand expand(MLOperand input, sequence<[EnforceRange] unsigned long> newShape, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand gather(MLOperand input, MLOperand indices, optional MLGatherOptions options = {});

  [RaisesException] MLOperand gatherElements(MLOperand input, MLOperand indices, optional MLGatherOptions options = {});

  [RaisesException] MLOperand gatherND(MLOperand input, MLOperand indices, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand gelu(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand gemm(MLOperand a, MLOperand b, optional MLGemmOptions options = {});

  [RaisesException] sequence<MLOperand> gru(MLOperand input, MLOperand weight, MLOperand recurrentWeight,
                    [EnforceRange] unsigned long steps, [EnforceRange] unsigned long hiddenSize,
                    optional MLGruOptions options = {});

  [RaisesException] MLOperand gruCell(MLOperand input, MLOperand weight, MLOperand recurrentWeight, MLOperand hiddenState,
                    [EnforceRange] unsigned long hiddenSize, optional MLGruCellOptions options = {});

  [RaisesException] MLOperand hardSigmoid(MLOperand x, optional MLHardSigmoidOptions options = {});

  [RaisesException] MLOperand hardSwish(MLOperand x, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand instanceNormalization(MLOperand input, optional MLInstanceNormalizationOptions options = {});

  [RaisesException] MLOperand matmul(MLOperand a, MLOperand b, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand layerNormalization(MLOperand input, optional MLLayerNormalizationOptions options = {});

  [RaisesException] MLOperand leakyRelu(MLOperand x, optional MLLeakyReluOptions options = {});

  [RaisesException] MLOperand linear(MLOperand input, optional MLLinearOptions options = {});

  [RaisesException] sequence<MLOperand> lstm(MLOperand input, MLOperand weight, MLOperand recurrentWeight,
                                             [EnforceRange] unsigned long steps, [EnforceRange] unsigned long hiddenSize,
                                             optional MLLstmOptions options = {});

  [RaisesException] sequence<MLOperand> lstmCell(MLOperand input, MLOperand weight, MLOperand recurrentWeight,
                                                 MLOperand hiddenState, MLOperand cellState, [EnforceRange] unsigned long hiddenSize,
                                                 optional MLLstmCellOptions options = {});

  [
    CallWith=ScriptState,
    RaisesException
  ] MLOperand pad(MLOperand input, sequence<[EnforceRange] unsigned long> beginningPadding,
                  sequence<[EnforceRange] unsigned long> endingPadding, optional MLPadOptions options = {});

  // Pooling operations
  [RaisesException] MLOperand averagePool2d(MLOperand input, optional MLPool2dOptions options = {});
  [RaisesException] MLOperand l2Pool2d(MLOperand input, optional MLPool2dOptions options = {});
  [RaisesException] MLOperand maxPool2d(MLOperand input, optional MLPool2dOptions options = {});

  [RaisesException] MLOperand prelu(MLOperand x, MLOperand slope, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand quantizeLinear(MLOperand input, MLOperand scale, MLOperand zeroPoint, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand reduceL1(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceL2(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceLogSum(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceLogSumExp(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceMax(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceMean(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceMin(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceProduct(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceSum(MLOperand input, optional MLReduceOptions options = {});
  [RaisesException] MLOperand reduceSumSquare(MLOperand input, optional MLReduceOptions options = {});

  [RaisesException] MLOperand relu(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand reshape(MLOperand input, sequence<[EnforceRange] unsigned long> newShape, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand reverse(MLOperand input, optional MLReverseOptions options = {});

  [
    CallWith=ScriptState,
    RaisesException
  ] MLOperand resample2d(MLOperand input, optional MLResample2dOptions options = {});

  [RaisesException] MLOperand scatterElements(MLOperand input, MLOperand indices, MLOperand updates, optional MLScatterOptions options = {});
  [RaisesException] MLOperand scatterND(MLOperand input, MLOperand indices, MLOperand updates, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand sigmoid(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand slice(MLOperand input,
                                    sequence<[EnforceRange] unsigned long> starts,
                                    sequence<[EnforceRange] unsigned long> sizes,
                                    optional MLSliceOptions options = {});

  // TODO: crbug.com/342919187 - Remove the deprecated version once ort-web 1.18.* supporting softmax(input, aixs) is released.
  [RaisesException] MLOperand softmax(MLOperand input, optional MLOperatorOptions options = {});
  [RaisesException] MLOperand softmax(MLOperand input, [EnforceRange] unsigned long axis, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand softplus(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand softsign(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] sequence<MLOperand> split(MLOperand input, [EnforceRange] unsigned long splits, optional MLSplitOptions options = {});
  [RaisesException] sequence<MLOperand> split(MLOperand input, sequence<[EnforceRange] unsigned long> splits, optional MLSplitOptions options = {});

  [RaisesException] MLOperand tanh(MLOperand input, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand tile(
      MLOperand input, sequence<[EnforceRange] unsigned long> repetitions, optional MLOperatorOptions options = {});

  [RaisesException] MLOperand transpose(
      MLOperand input, optional MLTransposeOptions options = {});

  [RaisesException] MLOperand triangular(MLOperand input, optional MLTriangularOptions options = {});

  [RaisesException] MLOperand where(MLOperand condition, MLOperand trueValue, MLOperand falseValue, optional MLOperatorOptions options = {});

  [
    CallWith=ScriptState,
    RaisesException
  ] Promise<MLGraph> build(MLNamedOperands outputs);
};