1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16//===----------------------------------------------------------------------===//
17//
18// This is the operation definition file for TensorFlow.js dialect operations.
19//
20//===----------------------------------------------------------------------===//
21
22#ifndef TFJS_DIALECT
23#define TFJS_DIALECT
24
25include "mlir/IR/OpBase.td"
26include "mlir/Interfaces/SideEffectInterfaces.td"
27
28//===----------------------------------------------------------------------===//
29// TensorFlow.js dialect definitions
30//===----------------------------------------------------------------------===//
31
32def TFJSDialect : Dialect {
33  let name = "tfjs";
34
35  let summary = "Types and operations for TensorFlow.js dialect";
36  let description = [{
37    This dialect contains operations for TensorFlow.js. This dialect will be
38    used in conjunction with the TensorFlow dialects for converting & optimizing
39    TF graphs to be deployed on TFJS.
40  }];
41
42  let cppNamespace = "::mlir::tfjs";
43}
44
45//===----------------------------------------------------------------------===//
46// TensorFlow.js op definitions
47//===----------------------------------------------------------------------===//
48
49// Base class for the operation in this dialect
50class TFJS_Op<string mnemonic, list<OpTrait> traits = []> :
51    Op<TFJSDialect, mnemonic, traits>;
52
53def TFJS_PReluOp : TFJS_Op<"Prelu", [NoSideEffect, ResultsBroadcastableShape,
54                                     SameOperandsAndResultElementType]> {
55  let summary = "Parametric Rectified Linear Unit operator";
56  let description = [{
57    Element-wise PReLU operator
58      x -> x >= 0 ? x : (alpha * x)
59  }];
60
61  let arguments = (ins AnyTensor:$input, AnyTensor:$alpha);
62  let results = (outs AnyTensor:$output);
63  let assemblyFormat =
64    " operands attr-dict `:` `(` type(operands) `)` `->` type($output)";
65}
66#endif // TFJS_DIALECT
67