From cc143617a75ddd4f86084ef41c91fbe5f3ccf051 Mon Sep 17 00:00:00 2001 From: Mehdi Amini Date: Wed, 17 Feb 2021 22:28:10 -0800 Subject: [PATCH] Ignore `use_locking` attribute when decomposing resources operations to target XLA This is the behavior of the existing bridge and this is aligning the implementation of the MLIR-based bridge with it. It also matches the original intent: the comment was explicit about the intent but the implementation wasn't matching it. PiperOrigin-RevId: 358107344 Change-Id: I789b5b97cd1dcab4099a7c0388d075d65cd13950 --- .../mlir/tensorflow/transforms/decompose_resource_ops.td | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td index 1b304bec9bc..0bfa1bec143 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td @@ -357,7 +357,7 @@ def DecomposeResourceApplyCenteredRMSProp : Pattern< (TF_ResourceApplyCenteredRMSPropOp:$src_op $var_resource, $mg_resource, $ms_resource, $mom_resource, $lr, $rho, $momentum, $epsilon, - $grad, ConstBoolAttrFalse:$use_locking + $grad, $use_locking ), [(TF_ConstOp:$one (GetScalarOfType<1> $grad)), (CreateTFReadVariableOp $src_op, $grad, $ms_resource), @@ -419,7 +419,7 @@ def DecomposeResourceApplyRMSProp : Pattern< (TF_ResourceApplyRMSPropOp:$src_op $var_resource, $ms_resource, $mom_resource, $lr, $rho, $momentum, $epsilon, - $grad, ConstBoolAttrFalse:$use_locking + $grad, $use_locking ), [(TF_ConstOp:$one (GetScalarOfType<1> $grad)), (CreateTFReadVariableOp $src_op, $grad, $ms_resource), @@ -456,7 +456,7 @@ def DecomposeResourceApplyProximalAdagrad : Pattern< (TF_ResourceApplyProximalAdagradOp:$src_op $var_resource, $accum_resource, $lr, $l1, $l2, $grad, - ConstBoolAttrFalse:$use_locking + $use_locking ), [(TF_ConstOp:$one (GetScalarOfType<1> $grad)), (TF_ConstOp:$zero (GetScalarOfType<0> $grad)),