experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Register TPUPartitionedOutput in TF MLIR ODS.

Browse Source 
PiperOrigin-RevId: 348503783
Change-Id: Ib918eb5657e2eb530082355bc46d08d0c7f0e131
This commit is contained in:
Andy Ly 2020-12-21 11:54:34 -08:00 committed by TensorFlower Gardener
parent 6ee022991d
commit 74cda1cdde
1 changed files with 24 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td
View File

@ -2062,4 +2062,28 @@ An op that groups a list of partitioned inputs together. This op
TF_DerivedOperandSizeAttr N = TF_DerivedOperandSizeAttr<0>;
}
def TF_TPUPartitionedOutputOp : TF_Op<"TPUPartitionedOutput", [NoSideEffect]> {
let summary = [{
An op that demultiplexes a tensor to be sharded by XLA to a list of partitioned
}];
let description = [{
outputs outside the XLA computation.
}];
let arguments = (ins
TF_Tensor:$inputs,
DefaultValuedAttr<I64Attr, "0">:$partition_dim,
OptionalAttr<StrAttr>:$_XlaSharding
);
let results = (outs
Variadic<TF_Tensor>:$output
);
TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
TF_DerivedResultSizeAttr num_splits = TF_DerivedResultSizeAttr<0>;
}
#endif // TF_OPS