Automated rollback of commit ed2a63c681b9f46913430a0c2a389ba3e63e3407
PiperOrigin-RevId: 263590924
diff --git a/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir b/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir
index fd35ed8..b77e7ad 100644
--- a/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir
@@ -321,6 +321,14 @@
// CHECK: %8 = "tf.MatMul"(%3, %7) {transpose_a = false, transpose_b = true} : (tensor<*xf32>, tensor<*xf32>) -> tensor<1x1000xf32>
}
+func @snapshot(%arg0: tensor<3xi32>) -> tensor<3xi32> {
+ %0 = "tf.Snapshot"(%arg0) : (tensor<3xi32>) -> tensor<3xi32>
+ return %0 : tensor<3xi32>
+ // Should be converted to Identity and then from Identity to value
+ // CHECK-LABEL: snapshot
+ // CHECK: return %arg0 : tensor<3xi32>
+}
+
func @stop_gradient(%arg0: tensor<3xi32>) -> tensor<3xi32> {
%0 = "tf.StopGradient"(%arg0) : (tensor<3xi32>) -> tensor<3xi32>
return %0 : tensor<3xi32>
diff --git a/tensorflow/compiler/mlir/lite/transforms/prepare_patterns.td b/tensorflow/compiler/mlir/lite/transforms/prepare_patterns.td
index 6b5b754..60b2acd 100644
--- a/tensorflow/compiler/mlir/lite/transforms/prepare_patterns.td
+++ b/tensorflow/compiler/mlir/lite/transforms/prepare_patterns.td
@@ -105,6 +105,7 @@
/*delta=*/(ConstantOp TFi32<-1>)), (ConstantOp TFi32<1>))), $b,
ConstBoolAttrFalse, $bt)>;
+def : Pat<(TF_SnapshotOp $arg), (TF_IdentityOp $arg)>;
def : Pat<(TF_StopGradientOp $arg), (TF_IdentityOp $arg)>;
//===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index e7d039d..b4e34fe 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -2848,47 +2848,9 @@
}
def TF_SelectV2Op : TF_Op<"SelectV2", [NoSideEffect]> {
- let summary = "Selects elements from `x` or `y`, depending on `condition`.";
+ let summary = "";
let description = [{
-The `x`, and `y` tensors must all have the same shape, and the
-output will also have that shape.
-
-The `condition` tensor must be a scalar if `x` and `y` are scalars.
-If `x` and `y` are vectors or higher rank, then `condition` must be either a
-scalar, a vector with size matching the first dimension of `x`, or must have
-the same shape as `x`.
-
-The `condition` tensor acts as a mask that chooses, based on the value at each
-element, whether the corresponding element / row in the output should be
-taken from `x` (if true) or `y` (if false).
-
-If `condition` is a vector and `x` and `y` are higher rank matrices, then
-it chooses which row (outer dimension) to copy from `x` and `y`.
-If `condition` has the same shape as `x` and `y`, then it chooses which
-element to copy from `x` and `y`.
-
-For example:
-
-```python
-# 'condition' tensor is [[True, False]
-# [False, True]]
-# 't' is [[1, 2],
-# [3, 4]]
-# 'e' is [[5, 6],
-# [7, 8]]
-select(condition, t, e) # => [[1, 6], [7, 4]]
-
-
-# 'condition' tensor is [True, False]
-# 't' is [[1, 2],
-# [3, 4]]
-# 'e' is [[5, 6],
-# [7, 8]]
-select(condition, t, e) ==> [[1, 2],
- [7, 8]]
-
-```
}];
let arguments = (ins
@@ -3005,6 +2967,23 @@
TF_DerivedOperandTypeAttr Index = TF_DerivedOperandTypeAttr<1>;
}
+def TF_SnapshotOp : TF_Op<"Snapshot", [NoSideEffect, SameOperandsAndResultType]> {
+ let summary = "Returns a copy of the input tensor.";
+
+ let description = [{
+ }];
+
+ let arguments = (ins
+ TF_Tensor:$input
+ );
+
+ let results = (outs
+ TF_Tensor:$output
+ );
+
+ TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
def TF_SoftmaxOp : TF_Op<"Softmax", [NoSideEffect]> {
let summary = "Computes softmax activations.";
