Support lstm and lstmCell operations

webmachinelearning · May 8, 2023 · 425f377 · 425f377
1 parent 4650499
commit 425f377
Show file tree

Hide file tree

Showing 5 changed files with 808 additions and 0 deletions.
diff --git a/src/nn/graph_builder.ts b/src/nn/graph_builder.ts
@@ -15,6 +15,7 @@ import {HardSigmoid} from './ops/hard_sigmoid';
 import {InstanceNormalization} from './ops/instance_norm';
 import {LeakyRelu} from './ops/leaky_relu';
 import {Linear} from './ops/linear';
+import {Lstm, LstmCell} from './ops/lstm';
 import {Pad} from './ops/pad';
 import {AveragePool2d, L2Pool2d, MaxPool2d} from './ops/pool2d';
 import {PRelu} from './ops/prelu';
@@ -211,9 +212,43 @@ export interface MLLinearOptions {
  beta?: number;
 }
 
+/**
+ * [spec](https://webmachinelearning.github.io/webnn/#enumdef-mllstmweightlayout)
+ */
+export enum MLLstmWeightLayout {
+ 'iofg' = 'iofg', // input-output-forget-cell gate ordering
+ 'ifgo' = 'ifgo' // input-forget-cell-output gate ordering
+}
+
 /**
  * [spec](https://webmachinelearning.github.io/webnn/#enumdef-mlpaddingmode)
  */
+export interface MLLstmOptions {
+ bias?: MLOperand;
+ recurrentBias?: MLOperand;
+ peepholeWeight?: MLOperand;
+ initialHiddenState?: MLOperand;
+ initialCellState?: MLOperand;
+ returnSequence?: boolean;
+ direction?: MLRecurrentNetworkDirection;
+ layout?: MLLstmWeightLayout;
+ activations?: MLActivation[];
+}
+
+/**
+ * [spec](https://webmachinelearning.github.io/webnn/#dictdef-mllstmcelloptions)
+ */
+export interface MLLstmCellOptions {
+ bias?: MLOperand;
+ recurrentBias?: MLOperand;
+ peepholeWeight?: MLOperand;
+ layout?: MLLstmWeightLayout;
+ activations?: MLActivation[];
+}
+
+/**
+ * [spec](https://webmachinelearning.github.io/webnn/#dictdef-mllstmoptions)
+ */
 export enum MLPaddingMode {
  'constant' = 'constant',
  'edge' = 'edge',
@@ -769,6 +804,38 @@ export class MLGraphBuilder {
  }
  }
 
+ /**
+ * [spec](https://webmachinelearning.github.io/webnn/#api-mlgraphbuilder-lstm)
+ */
+ lstm(input: MLOperand, weight: MLOperand, recurrentWeight: MLOperand,
+ steps: number, hiddenSize: number, options: MLLstmOptions = {}):
+ MLOperand[] {
+ this.validateOperandBuilder([
+ input, weight, recurrentWeight, options.bias, options.recurrentBias,
+ options.peepholeWeight, options.initialHiddenState,
+ options.initialCellState
+ ]);
+ return (new Lstm(
+ input, weight, recurrentWeight, steps, hiddenSize, options))
+ .outputs;
+ }
+
+ /**
+ * [spec](https://webmachinelearning.github.io/webnn/#api-mlgraphbuilder-lstmcell)
+ */
+ lstmCell(input: MLOperand, weight: MLOperand, recurrentWeight: MLOperand,
+ hiddenState: MLOperand, cellState: MLOperand, hiddenSize: number,
+ options: MLLstmCellOptions = {}): MLOperand[] {
+ this.validateOperandBuilder([
+ input, weight, recurrentWeight, hiddenState, cellState, options.bias,
+ options.recurrentBias, options.peepholeWeight
+ ]);
+ return (new LstmCell(
+ input, weight, recurrentWeight, hiddenState, cellState,
+ hiddenSize, options))
+ .outputs;
+ }
+
  /**
  * [spec](https://webmachinelearning.github.io/webnn/#dom-mlgraphbuilder-matmul)
  */