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

Fix decomposition of ResourceApplyCenteredRMSProp op pattern.

Browse Source 
AssignSubVariableOp only needs the resource handle and the tensor to be subtracted from it.

PiperOrigin-RevId: 316597384
Change-Id: I26a99cb6fe210cc66429ebca13765af3109748ae
This commit is contained in:
Prakalp Srivastava 2020-06-15 19:50:34 -07:00 committed by TensorFlower Gardener
parent 552604a2ef
commit c118bc4b5e
2 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
View File

@ -411,6 +411,7 @@ func @decompose_resource_apply_centered_RMS_prop(%arg0: tensor<f32>, %arg1: tens
// CHECK: [[VAR:%.*]] = "tf.ReadVariableOp"([[VAR_HANDLE]])
// CHECK: [[VAR_NEW:%.*]] = "tf.Sub"([[VAR]], [[MOM_NEW]])
// CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[VAR_NEW]])
"tf.ResourceApplyCenteredRMSProp"(%0, %1, %2, %3, %arg4, %arg5, %arg6, %arg7, %arg8) {use_locking = false} : (tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>) -> ()
return

5
tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
View File

@ -387,9 +387,6 @@ def DecomposeResourceApplyCenteredRMSProp :
),
(TF_AssignVariableOp $mom_resource, $mom_new),
// var <- var - mom
(TF_AssignSubVariableOp $var_resource,
(TF_SubOp (CreateTFReadVariableOp $src_op, $grad, $var_resource),
$mom_new)
)
(TF_AssignSubVariableOp $var_resource, $mom_new)
]
>;